22 lines
567 B
Python
22 lines
567 B
Python
from typing import Union
|
|
|
|
COMPLETE_MODELS = {"j2-light": 8191, "j2-mid": 8191, "j2-ultra": 8191}
|
|
|
|
|
|
def ai21_model_to_context_size(model: str) -> Union[int, None]:
|
|
"""Calculate the maximum number of tokens possible to generate for a model.
|
|
|
|
Args:
|
|
model: The modelname we want to know the context size for.
|
|
|
|
Returns:
|
|
The maximum context size
|
|
|
|
"""
|
|
token_limit = COMPLETE_MODELS.get(model, None)
|
|
|
|
if token_limit is None:
|
|
raise ValueError(f"Model name {model} not found in {COMPLETE_MODELS.keys()}")
|
|
|
|
return token_limit
|