llama-index

1
from typing import Dict
2

3
LLAMA_MODELS = {
4
    "meta-llama/Llama-2-7b-chat-hf": 4096,
5
}
6

7
ALL_AVAILABLE_MODELS = {
8
    **LLAMA_MODELS,
9
}
10

11
DISCONTINUED_MODELS: Dict[str, int] = {}
12

13

14
def everlyai_modelname_to_contextsize(modelname: str) -> int:
15
    """Calculate the maximum number of tokens possible to generate for a model.
16

17
    Args:
18
        modelname: The modelname we want to know the context size for.
19

20
    Returns:
21
        The maximum context size
22

23
    Example:
24
        .. code-block:: python
25

26
            max_tokens = everlyai_modelname_to_contextsize(model_name)
27
    """
28
    if modelname in DISCONTINUED_MODELS:
29
        raise ValueError(
30
            f"EverlyAI hosted model {modelname} has been discontinued. "
31
            "Please choose another model."
32
        )
33

34
    context_size = ALL_AVAILABLE_MODELS.get(modelname, None)
35

36
    if context_size is None:
37
        raise ValueError(
38
            f"Unknown model: {modelname}. Please provide a valid EverlyAI model name."
39
            "Known models are: " + ", ".join(ALL_AVAILABLE_MODELS.keys())
40
        )
41

42
    return context_size
43