Update app.py

app.py

@@ -5,14 +5,18 @@ import os
 import re
 import emoji
 
-MODEL_NAME = "NeuroSpaceX/ruSpamNS"
 TOKEN = os.getenv("HF_TOKEN")
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=TOKEN)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, use_auth_token=TOKEN)
+models = {
+    "ruSpamNS_v13": "NeuroSpaceX/ruSpamNS_v13",
+    "ruSpamNS_big": "NeuroSpaceX/ruSpamNS_big",
+    "ruSpamNS_small": "NeuroSpaceX/ruSpamNS_small"
+}
+
+tokenizers = {name: AutoTokenizer.from_pretrained(path, use_auth_token=TOKEN) for name, path in models.items()}
+models = {name: AutoModelForSequenceClassification.from_pretrained(path, use_auth_token=TOKEN).to(torch.device('cuda' if torch.cuda.is_available() else 'cpu')) for name, path in models.items()}
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model.to(device)
 
 def clean_text(text):
     text = emoji.replace_emoji(text, replace='')
@@ -23,18 +27,18 @@ def clean_text(text):
     return text
 
 def classify_text(text, model_choice):
-    model_name = f"NeuroSpaceX/ruSpamNS_{model_choice}"
-    tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=TOKEN)
-    model = AutoModelForSequenceClassification.from_pretrained(model_name, use_auth_token=TOKEN)
-    model.to(device)
+    tokenizer = tokenizers[model_choice]
+    model = models[model_choice]
     
     message = clean_text(text)
     encoding = tokenizer(message, padding='max_length', truncation=True, max_length=128, return_tensors='pt')
     input_ids = encoding['input_ids'].to(device)
     attention_mask = encoding['attention_mask'].to(device)
+    
     with torch.no_grad():
         outputs = model(input_ids, attention_mask=attention_mask).logits
         prediction = torch.sigmoid(outputs).cpu().numpy()[0][0]
+    
     label = "СПАМ" if prediction >= 0.5 else "НЕ СПАМ"
     return f"{label} (вероятность: {prediction*100:.2f}%)"
 
@@ -42,7 +46,7 @@ iface = gr.Interface(
     fn=classify_text,
     inputs=[
         gr.Textbox(lines=3, placeholder="Введите текст..."),
-        gr.Radio(["small", "big"], label="Выберите модель")
+        gr.Radio(["ruSpamNS_v13", "ruSpamNS_big", "ruSpamNS_small"], label="Выберите модель")
     ],
     outputs="text",
     title="ruSpamNS - Проверка на спам",