translations/model_name_mapping.py

from dataclasses import dataclass, field
from typing import Dict, Literal, Tuple
from pydantic import BaseModel, ValidationError
from enum import Enum


class ModelNameMapping(BaseModel):
    data: Dict[Tuple, str] = field(default_factory=dict)

    def __getitem__(self, key: Tuple) -> str:
        return self.data[key]

    def get(self, key: Tuple) -> str:
        return self.data.get(key, "")


class LanguagesEnum(Enum):
    es = "es"
    en = "en"
    ca = "ca"
    pt = "pt"


es_to_en: Tuple = (LanguagesEnum.es.value, LanguagesEnum.en.value)
es_to_ca: Tuple = (LanguagesEnum.es.value, LanguagesEnum.ca.value)
es_to_po: Tuple = (LanguagesEnum.es.value, LanguagesEnum.pt.value)

en_to_es: Tuple = (LanguagesEnum.en.value, LanguagesEnum.es.value)
en_to_po: Tuple = (LanguagesEnum.en.value, LanguagesEnum.pt.value)

ca_to_es: Tuple = (LanguagesEnum.ca.value, LanguagesEnum.es.value)

po_to_en: Tuple = (LanguagesEnum.pt.value, LanguagesEnum.en.value)
po_to_es: Tuple = (LanguagesEnum.pt.value, LanguagesEnum.es.value)

models = {
    es_to_en: "Helsinki-NLP/opus-mt-es-en",
    es_to_ca: "Helsinki-NLP/opus-mt-es-ca",
    en_to_es: "Helsinki-NLP/opus-mt-en-es",
    en_to_po: "Helsinki-NLP/opus-mt-tc-big-en-pt",
    ca_to_es: "Helsinki-NLP/opus-mt-ca-es",
    po_to_en: "Helsinki-NLP/opus-mt-pt-en",
    po_to_es: "Helsinki-NLP/opus-mt-pt-es",
}


MODEL_NAME_MAPPING = ModelNameMapping(data=models)