Update app.py

app.py

@@ -1,11 +1,17 @@
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import os
 
 def load_model(model_name):
+    token = os.getenv("HG_TOKEN")
+    if not token:
+        raise ValueError("Hugging Face API token not found. Please set HG_TOKEN environment variable.")
+    
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=1).to(device).eval()
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    # Load model with authentication token if necessary
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=1, use_auth_token=token).to(device).eval()
+    tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=token)
     return model, tokenizer, device
 
 def classify_message(model, tokenizer, device, message, model_name):
@@ -26,7 +32,12 @@ def classify_message(model, tokenizer, device, message, model_name):
     return "Спам" if is_spam else "Не спам"
 
 def spam_classifier_interface(message, model_choice):
-    model_name = 'NeuroSpaceX/ruSpamNS_v1' if model_choice == "Model v1" else 'NeuroSpaceX/ruSpamNS_v6'
+    # Choose model based on user's choice
+    model_name = {
+        "Model v1": 'NeuroSpaceX/ruSpamNS_v1',
+        "Model v6": 'NeuroSpaceX/ruSpamNS_v6',
+        "Model v7beta": 'NeuroSpaceX/ruSpamNS_v7'
+    }[model_choice]
     model, tokenizer, device = load_model(model_name)
     return classify_message(model, tokenizer, device, message, model_name)
 
@@ -35,7 +46,7 @@ interface = gr.Interface(
     fn=spam_classifier_interface,
     inputs=[
         gr.Textbox(label="Введите сообщение для классификации", placeholder="Введите текст..."),
-        gr.Radio(["Model v1", "Model v6"], label="Выберите модель")
+        gr.Radio(["Model v1", "Model v6", "Model v7beta"], label="Выберите модель")
     ],
     outputs=gr.Textbox(label="Результат"),
     title="Классификатор Спам/Не Спам",