File size: 1,128 Bytes
c7256ee | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | from typing import Any
from fastapi import HTTPException
from models.llama_3_8b import Llama3_8B
from models.mistral_7b import Mistral_7b
from models.qwen_2_5 import Qwen2_5
from models.deepseek_v3 import DeepSeek_V3
from models.tiny_aya import TinyAya
# model defination
# copied from /models
def build_models(hf_token: str) -> dict[str, Any]:
return {
"Llama-3-8B": Llama3_8B(token=hf_token),
"Mistral-7B": Mistral_7b(token=hf_token),
"Qwen-2.5": Qwen2_5(token=hf_token),
"DeepSeek-V3": DeepSeek_V3(token=hf_token),
"TinyAya": TinyAya(token=hf_token),
}
def resolve_model(name: str, models: dict[str, Any]) -> tuple[str, Any]:
aliases = {
"llama": "Llama-3-8B",
"mistral": "Mistral-7B",
"qwen": "Qwen-2.5",
"deepseek": "DeepSeek-V3",
"tinyaya": "TinyAya",
}
model_key = aliases.get(name.lower(), name)
if model_key not in models:
allowed = ", ".join(models.keys())
raise HTTPException(status_code=400, detail=f"Unknown model '{name}'. Use one of: {allowed}")
return model_key, models[model_key]
|