Better logic to determine context window size

This commit is contained in:
Paul Gauthier 2023-06-21 14:10:30 -07:00
parent 9d1c53489e
commit 6e3738c901

View file

@ -1,15 +1,27 @@
import re
known_tokens = {
"gpt-3.5-turbo": 4,
"gpt-4": 8,
}
class Model:
def __init__(self, name, tokens=None):
def __init__(self, name):
self.name = name
if tokens is None:
tokens = None
match = re.search(r"-([0-9]+)k", name)
if match:
tokens = int(match.group(1))
else:
for m, t in known_tokens.items():
if name.startswith(m):
tokens = t
default_tokens = 8
tokens = int(match.group(1)) if match else default_tokens
if tokens is None:
raise ValueError(f"Unknown context window size for model: {name}")
self.max_context_tokens = tokens * 1024
@ -28,6 +40,6 @@ class Model:
return self.is_gpt35()
GPT4 = Model("gpt-4", 8)
GPT4 = Model("gpt-4")
GPT35 = Model("gpt-3.5-turbo")
GPT35_16k = Model("gpt-3.5-turbo-16k")