"""
model_contexts.py

Contains a reference dictionary that defines the maximum context lengths for various
large language models (LLMs) from different providers.
Each model is represented as a key in the dictionary `MODEL_CONTEXT_LENGTHS`, with the associated maximum context
length as the value. These context lengths indicate the maximum token capacity that each model can handle within a
single input prompt.

The `MODEL_CONTEXT_LENGTHS` dictionary includes models from the following providers:
- **OpenAI** (e.g., GPT-3.5, GPT-4 variants)
- **Anthropic** (e.g., Claude models)
- **Google** (e.g., Gemini, PaLM)
- **Meta** (LLaMA models)
- **Cohere** (e.g., Command models)
- **AI21 Labs** (e.g., Jurassic-2)
- **EleutherAI** (e.g., GPT-J, GPT-NeoX)
- **Stability AI** and **Microsoft** (e.g., StableLM, Orca)
"""

MODEL_CONTEXT_LENGTHS = {
    # OpenAI models
    "gpt-3.5-turbo": 16_385,
    "gpt-3.5-turbo-0125": 16_385,
    "gpt-3.5-turbo-0301": 4096,
    "gpt-3.5-turbo-0613": 4096,
    "gpt-3.5-turbo-1106": 16_385,
    "gpt-3.5-turbo-16k": 16_385,
    "gpt-4": 8192,
    "gpt-4-0125-preview": 128_000,
    "gpt-4-0314": 8192,
    "gpt-4-0613": 8192,
    "gpt-4-1106-preview": 128_000,
    "gpt-4-1106-vision-preview": 128_000,
    "gpt-4-32k": 32_768,
    "gpt-4-32k-0314": 32_768,
    "gpt-4-32k-0613": 32_768,
    "gpt-4-turbo-2024-04-09": 128_000,
    "gpt-4-turbo-preview": 128_000,
    "gpt-4-vision-preview": 128_000,
    "gpt-4o": 128_000,
    "gpt-4o-2024-11-20": 128_000,
    "gpt-4o-2024-08-06": 128_000,
    "gpt-4o-2024-05-13": 128_000,
    "gpt-4o-mini": 128_000,
    "gpt-4o-mini-2024-07-18": 128_000,
    "o1": 200_000,
    "o1-2024-12-17": 200_000,
    "o1-mini": 128_000,
    "o1-mini-2024-09-12": 128_000,
    "o1-preview": 128_000,
    "o1-preview-2024-09-12": 128_000,
    # Anthropic models
    "claude-1": 9000,
    "claude-1.3": 100_000,
    "claude-instant-1": 9000,
    "claude-2.0": 100_000,
    "claude-2.1": 200_000,
    "claude-instant-1.2": 100_000,
    "claude-3-opus-20240229": 200_000,
    "claude-3-sonnet-20240229": 200_000,
    "claude-3-haiku-20240307": 200_000,
    # Google models
    "palm": 8192,
    "gemini-pro": 32_768,
    "gemini-ultra": 32_768,
    # Meta models (LLaMA 2)
    "llama-2-7b": 4096,
    "llama-2-13b": 4096,
    "llama-2-70b": 4096,
    # Cohere models
    "command": 4096,
    "command-light": 4096,
    "command-nightly": 4096,
    # AI21 Labs models
    "jurassic-2-ultra": 8192,
    "jurassic-2-mid": 8192,
    "jurassic-2-light": 8192,
    # EleutherAI models
    "gpt-j-6b": 2048,
    "gpt-neox-20b": 2048,
    # Stability AI models
    "stablelm-tuned-alpha-7b": 4096,
    # Microsoft models
    "orca": 2048,
}
