llama-index

1
from typing import Union
2

3
COMPLETE_MODELS = {"j2-light": 8191, "j2-mid": 8191, "j2-ultra": 8191}
4

5

6
def ai21_model_to_context_size(model: str) -> Union[int, None]:
7
    """Calculate the maximum number of tokens possible to generate for a model.
8

9
    Args:
10
        model: The modelname we want to know the context size for.
11

12
    Returns:
13
        The maximum context size
14

15
    """
16
    token_limit = COMPLETE_MODELS.get(model, None)
17

18
    if token_limit is None:
19
        raise ValueError(f"Model name {model} not found in {COMPLETE_MODELS.keys()}")
20

21
    return token_limit
22