Update app.py

app.py

@@ -1,9 +1,11 @@
+import os
 import time
 import torch
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel
 
+HF_TOKEN = os.getenv("HF_TOKEN")
 MODEL_ID = "ruSpamModels/ruSpam-Qwen-0.5B-50k"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -13,12 +15,20 @@ base_model = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.float16 if device == "cuda" else torch.float32,
     device_map=device,
     trust_remote_code=True,
+    token=HF_TOKEN,
 )
 
-model = PeftModel.from_pretrained(base_model, MODEL_ID)
+model = PeftModel.from_pretrained(
+    base_model,
+    MODEL_ID,
+    token=HF_TOKEN,
+)
 model.eval()
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID,
+    token=HF_TOKEN,
+)
 
 def classify(message):
     prompt = (
@@ -42,7 +52,6 @@ def classify(message):
         )
 
     elapsed = (time.time() - start) * 1000
-
     new_token_id = out[0, inputs["input_ids"].shape[1]]
     answer = tokenizer.decode(new_token_id).strip().lower()
 
@@ -57,11 +66,9 @@ def classify(message):
 
 iface = gr.Interface(
     fn=classify,
-    inputs=gr.Textbox(lines=4, placeholder="Введите сообщение"),
-    outputs=gr.Textbox(label="Результат"),
+    inputs=gr.Textbox(lines=4),
+    outputs=gr.Textbox(),
     title="ruSpam Qwen 0.5B",
-    description="Классификация сообщений: SPAM / HAM",
 )
 
-if __name__ == "__main__":
-    iface.launch()
+iface.launch()