Upload 5 files

model_wrapper/bert_wrapper.py

@@ -0,0 +1,37 @@
+import torch
+from transformers import AutoModelForSequenceClassification
+from transformers import BertTokenizerFast
+
+
+class BertWrapper(object):
+    MODELS_DIR: str = "new_models/"
+    MODEL_NAME: str = "model"
+    TOKENIZER: str = "tokenizer"
+
+    def __init__(self) -> None:
+        self.model = AutoModelForSequenceClassification.from_pretrained(
+            BertWrapper.MODELS_DIR + BertWrapper.MODEL_NAME, torchscript=True
+        )
+        self.tokenizer = BertTokenizerFast.from_pretrained(
+            BertWrapper.MODELS_DIR + BertWrapper.TOKENIZER
+        )
+        self.id2label: dict[int, str] = {0: "__label__positive", 1: "__label__negative"}
+
+    @torch.no_grad()
+    def __call__(self, text: str) -> str:
+        max_input_length = (
+            self.model.config.max_position_embeddings
+        )  # 512 for this model
+        inputs = self.tokenizer(
+            text,
+            max_length=max_input_length,
+            padding=True,
+            truncation=True,
+            return_tensors="pt",
+        )
+        outputs = self.model(
+            **inputs, return_dict=True
+        )  # output is logits for huggingfcae transformers
+        predicted = torch.nn.functional.softmax(outputs.logits, dim=1)
+        predicted_id = torch.argmax(predicted, dim=1).numpy()[0]
+        return self.id2label[predicted_id]

model_wrapper/fasttext_wrapper.py

@@ -0,0 +1,10 @@
+import fasttext
+
+class FasttextWrapper(object):
+    MODEL_PATH: str = "./model_fasttext.bin"
+
+    def __init__(self) -> None:
+        self.model = fasttext.load_model(FasttextWrapper.MODEL_PATH)
+
+    def __call__(self, text: str) -> str:
+        return self.model.predict(text)

model_wrapper/frida_wrapper.py

@@ -0,0 +1,8 @@
+# Захреначьте импорты сюда и сделайте метод call, как в других обёртках
+
+class FridaWrapper(object):
+    def __init__(self) -> None:
+        pass
+    
+    def __call__(self, text: str) -> str:
+        pass

model_wrapper/model_wrapper.py

@@ -0,0 +1,16 @@
+from bert_wrapper import BertWrapper
+from fasttext_wrapper import FasttextWrapper
+from frida_wrapper import FridaWrapper
+
+from typing import Any
+
+class ModelWrapper(object):
+    def __init__(self) -> None:
+        self.models_dict: dict[str, Any] = {
+            "fasttext": FasttextWrapper(),
+            "ru-BERT": BertWrapper(),
+            "FRIDA": FridaWrapper(),
+        }
+
+    def __call__(self, text: str, model_name: str) -> str:
+        return self.models_dict[model_name](text)