Add translation models

category_classification/models/models.py

@@ -6,6 +6,8 @@ import warnings
 from pathlib import Path
 
 from . import pipeline
+from .translation import create_translation_models
+
 
 def import_model_module(file_path: os.PathLike):
     module_name = str(Path(file_path).relative_to(os.getcwd())).replace(
@@ -76,6 +78,10 @@ for path in file_dir.glob("*"):
             language_to_models.setdefault(lang, {})
         language_to_models[lang][name] = get_model
 
+translation_models = create_translation_models(language_to_models["en"])
+language_to_models.setdefault("ru", {}).update(translation_models)
+models.update(translation_models)
+
 
 def get_model(name: str):
     if name not in models:
@@ -86,6 +92,7 @@ def get_model(name: str):
 def get_all_model_names():
     return list(models.keys())
 
+
 def get_model_names_by_lang(lang):
     if lang not in language_to_models:
         return []

category_classification/models/pipeline.py

@@ -1,4 +1,5 @@
 import typing as tp
+from collections import namedtuple
 
 import torch
 

category_classification/models/translation.py

@@ -0,0 +1,46 @@
+from collections import namedtuple
+from functools import partial
+
+import torch
+from transformers import pipeline
+
+
+def get_translator():
+    return pipeline(
+        "translation_en_to_ru",
+        model="Helsinki-NLP/opus-mt-ru-en",
+        device="cuda" if torch.cuda.is_available() else "cpu",
+        torch_dtype="auto",
+    )
+
+class Input:
+    def __init__(self, title, abstract, authors):
+        self.title = title
+        self.abstract = abstract
+        self.authors = authors
+    
+class TranslationModel:
+    def __init__(self, get_model):
+        self.translator = get_translator()
+        self.model = get_model()
+
+    def __call__(self, input):
+        def translate(text):
+            if text is None or text.strip() == "":
+                return ""
+            text = str(text).strip()
+            translated = self.translator(text)[0]['translation_text']
+            return translated
+        title = translate(input.title)
+        abstract = translate(input.abstract)
+        authors = translate(input.authors)
+        out = self.model(Input(title, abstract, authors))
+        return out
+
+
+def create_translation_models(models):
+    return {
+        f"{name} (С помощью перевода)": partial(TranslationModel, get_model=get_model)
+        for name, get_model in models.items()
+    }
+

requirements.txt

@@ -130,6 +130,7 @@ safetensors==0.5.3
 scikit-learn==1.6.1
 scipy==1.15.2
 Send2Trash==1.8.3
+sentencepiece==0.2.0
 six==1.17.0
 smmap==5.0.2
 sniffio==1.3.1