Update app.py

app.py

@@ -2,6 +2,9 @@ import streamlit as st
 import base64
 import fasttext
 import re
+import torch
+from transformers import AutoModelForSequenceClassification
+from transformers import BertTokenizerFast
 
 st.set_page_config(
     page_title="detoxi.ai",
@@ -54,6 +57,40 @@ st.markdown(
 st.write("""<p style='text-align: center; font-size: 24px;'>Это приложение сделает твою речь менее токсичной. 
 И даже не придётся платить 300 bucks.</p>""", unsafe_allow_html=True)
 
+class ModelWrapper(object):
+    MODELS_DIR: str = "./new_models/"
+    MODEL_NAME: str = "model"
+    TOKENIZER: str = "tokenizer"
+
+    def __init__(self):
+        self.model = AutoModelForSequenceClassification.from_pretrained(
+            ModelWrapper.MODELS_DIR + ModelWrapper.MODEL_NAME, torchscript=True
+        )
+        self.tokenizer = BertTokenizerFast.from_pretrained(
+            ModelWrapper.MODELS_DIR + ModelWrapper.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]
+
+
 def highlight_obscene_words(text):
     label,_=model.predict(text.lower())
     if label[0]=='__label__positive':