senga-dnotes / config.py
serenarolloh's picture
Update config.py
2468ba3 verified
raw
history blame
573 Bytes
from pydantic import BaseSettings
from typing import Dict
class Settings(BaseSettings):
# A dictionary to map shipper_id to the appropriate model
model_map: Dict[int, str] = {
61: "donut-v16",
81: "donut-v16",
139: "donut-v16",
165: "donut-v17",
145: "donut-v17",
127: "donut-v17",
}
space_base: str = "senga-ml"
def get_model_url(self, shipper_id: int) -> str:
model = self.model_map.get(shipper_id, "default-model")
return f"https://{self.space_base}/{model}"
settings = Settings()