use modified config

tasks/text.py

@@ -53,18 +53,36 @@ async def evaluate_text(request: TextEvaluationRequest):
     # MODEL INFERENCE CODE
     #--------------------------------------------------------------------------------------------   
     
-    
     try:
         # Set device
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         
         # Model and tokenizer paths
-        model_name = "Tonic/climate-guard-toxic-agent"
+        model_name = "answerdotai/ModernBERT-base"
+        
+        # Load config and modify it
+        config = AutoConfig.from_pretrained(model_name)
+        config.num_labels = 8
+        config.id2label = {
+            "0": "0_not_relevant",
+            "1": "1_not_happening",
+            "2": "2_not_human",
+            "3": "3_not_bad",
+            "4": "4_solutions_harmful_unnecessary",
+            "5": "5_science_is_unreliable",
+            "6": "6_proponents_biased",
+            "7": "7_fossil_fuels_needed"
+        }
+        config.label2id = {v: int(k) for k, v in config.id2label.items()}
+        config.problem_type = "single_label_classification"
         
-        # Load tokenizer and model
+        # Load tokenizer
         tokenizer = AutoTokenizer.from_pretrained(model_name)
+        
+        # Load model with modified config
         model = AutoModelForSequenceClassification.from_pretrained(
-            model_name,
+            "Tonic/climate-guard-toxic-agent",
+            config=config,
             torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
             trust_remote_code=True
         ).to(device)
@@ -123,7 +141,6 @@ async def evaluate_text(request: TextEvaluationRequest):
         print(f"Error during model inference: {str(e)}")
         raise
 
-
     #--------------------------------------------------------------------------------------------
     # MODEL INFERENCE ENDS HERE
     #--------------------------------------------------------------------------------------------