Spaces:
Sleeping
Sleeping
File size: 1,407 Bytes
10af882 |
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 39 40 41 42 43 44 45 46 47 48 |
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)
|