Update tasks/text.py

tasks/text.py

@@ -65,9 +65,14 @@ async def evaluate_text(request: TextEvaluationRequest):
 
     # Load model and tokenizer from Hugging Face Hub
     MODEL_REPO = "ClimateDebunk/FineTunedDistilBert4SeqClass"
+    MODEL_PATH = "./distilbert_trained.pth"
     tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
-    model = AutoModelForSequenceClassification.from_pretrained(MODEL_REPO)
-    model.eval()  # Set to evaluation mode
+
+    config = DistilBertConfig.from_pretrained("distilbert-base-uncased", num_labels=8)
+    model = DistilBertForSequenceClassification(config)
+    model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device("cpu")))
+    #model = AutoModelForSequenceClassification.from_pretrained(MODEL_REPO)
+    model.eval()  
     
     def preprocess(texts):
         """ Tokenize text inputs for DistilBERT """