Update app.py

app.py

@@ -1,10 +1,14 @@
 import gradio as gr
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import torch
+import os
+
+token = os.getenv('token')
 
 model_name = 'synthome-fyi-paper-classification'
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+model = AutoModelForSequenceClassification.from_pretrained(model_name, use_auth_token=token)
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=token)
 
 def classify_text(text):
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
@@ -12,7 +16,8 @@ def classify_text(text):
         outputs = model(**inputs)
     logits = outputs.logits
     prediction = logits.argmax().item()
-    return prediction
+    labels = ['AI only', 'Bio only', 'AIxBio']
+    return labels[prediction]
 
 iface = gr.Interface(fn=classify_text, inputs="text", outputs="text")
 iface.launch()