Update app.py

app.py

@@ -4,13 +4,19 @@ import torch.serialization
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
 @st.cache_resource
-def load_model():   
-    checkpoint = torch.load('TinyBERT_cls_model.pt', 
+def load_model():
+    #trained_model = 'TinyBERT_cls_model.pt'
+    #base_model = 'huawei-noah/TinyBERT_General_4L_312D'
+    
+    trained_model = 'distilbert-base_cls_model.pt'
+    base_model = 'distilbert-base-uncased'
+    
+    checkpoint = torch.load(trained_model, 
                           map_location='cpu',
                           weights_only=False)  
     
     model = AutoModelForSequenceClassification.from_pretrained(
-        "huawei-noah/TinyBERT_General_4L_312D",
+        base_model,
         num_labels=len(checkpoint['idx_to_category'])
     )
     model.load_state_dict(checkpoint['model_state_dict'])
@@ -20,7 +26,7 @@ def load_model():
     
     return model, tokenizer, idx_to_category
 
-def predict(title, abstract, model, tokenizer, idx_to_category, threshold=0.95):
+def predict(title, abstract, model, tokenizer, idx_to_category, threshold=0.97):
     text = f"{title} /n {abstract}" if abstract else title
     inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)