llama-index

1
import os
2
from typing import Optional, Union
3

4
WATSONX_MODELS = {
5
    "google/flan-t5-xxl": 4096,
6
    "google/flan-ul2": 4096,
7
    "bigscience/mt0-xxl": 4096,
8
    "eleutherai/gpt-neox-20b": 8192,
9
    "bigcode/starcoder": 8192,
10
    "meta-llama/llama-2-70b-chat": 4096,
11
    "ibm/mpt-7b-instruct2": 2048,
12
    "ibm/granite-13b-instruct-v1": 8192,
13
    "ibm/granite-13b-chat-v1": 8192,
14
}
15

16

17
def watsonx_model_to_context_size(model_id: str) -> Union[int, None]:
18
    """Calculate the maximum number of tokens possible to generate for a model.
19

20
    Args:
21
        model_id: The model name we want to know the context size for.
22

23
    Returns:
24
        The maximum context size
25
    """
26
    token_limit = WATSONX_MODELS.get(model_id, None)
27

28
    if token_limit is None:
29
        raise ValueError(f"Model name {model_id} not found in {WATSONX_MODELS.keys()}")
30

31
    return token_limit
32

33

34
def get_from_param_or_env_without_error(
35
    param: Optional[str] = None,
36
    env_key: Optional[str] = None,
37
) -> Union[str, None]:
38
    """Get a value from a param or an environment variable without error."""
39
    if param is not None:
40
        return param
41
    elif env_key and env_key in os.environ and os.environ[env_key]:
42
        return os.environ[env_key]
43
    else:
44
        return None
45