revert to reference code

tasks/text.py

@@ -1,4 +1,3 @@
-
 from fastapi import APIRouter
 from datetime import datetime
 import time
@@ -14,7 +13,7 @@ from huggingface_hub import login
 from dotenv import load_dotenv
 
 from .utils.evaluation import TextEvaluationRequest
-from .utils.emissions import tracker, clean_emissions_data, get_space_info, start_tracking, stop_tracking
+from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 # Load environment variables
 load_dotenv()
@@ -29,7 +28,7 @@ os.environ["TORCH_COMPILE_DISABLE"] = "1"
 
 router = APIRouter()
 
-DESCRIPTION = "Climate Guard Toxic Agent Classifier"
+DESCRIPTION = "ModernBERT Climate Claims Classifier"
 ROUTE = "/text"
 
 class TextClassifier:
@@ -43,13 +42,15 @@ class TextClassifier:
                 # Load config
                 self.config = AutoConfig.from_pretrained(
                     model_name,
+                    num_labels=8,
+                    problem_type="single_label_classification",
                     trust_remote_code=True
                 )
                 
                 # Initialize tokenizer
                 self.tokenizer = AutoTokenizer.from_pretrained(
                     model_name,
-                    model_max_length=2048,
+                    model_max_length=8192,
                     padding_side='right',
                     truncation_side='right',
                     trust_remote_code=True
@@ -60,15 +61,11 @@ class TextClassifier:
                     model_name,
                     config=self.config,
                     trust_remote_code=True,
-                    torch_dtype=torch.float32,
-                    device_map="auto",
-                    low_cpu_mem_usage=True
+                    torch_dtype=torch.float32
                 )
                 
-                # Force model to CPU if CUDA is not available
-                if not torch.cuda.is_available():
-                    self.model = self.model.cpu()
-                
+                # Move model to appropriate device
+                self.model = self.model.to(self.device)
                 self.model.eval()
                 print("Model initialized successfully")
                 break
@@ -84,12 +81,12 @@ class TextClassifier:
         try:
             print(f"Processing batch {batch_idx} with {len(batch)} items")
             
-            # Tokenize with smaller max length
+            # Tokenize
             inputs = self.tokenizer(
                 batch,
                 return_tensors="pt",
                 truncation=True,
-                max_length=512,  # Reduced max length
+                max_length=512,
                 padding=True
             )
             
@@ -129,14 +126,14 @@ async def evaluate_text(request: TextEvaluationRequest):
         "2_not_human": 2,
         "3_not_bad": 3,
         "4_solutions_harmful_unnecessary": 4,
-        "5_science_unreliable": 5,
+        "5_science_is_unreliable": 5,
         "6_proponents_biased": 6,
         "7_fossil_fuels_needed": 7
     }
 
     try:
         # Load and prepare the dataset
-        dataset = load_dataset("QuotaClimat/frugalaichallenge-text-train", token=HF_TOKEN)
+        dataset = load_dataset(request.dataset_name)
         
         # Convert string labels to integers
         dataset = dataset.map(lambda x: {"label": LABEL_MAPPING[x["label"]]})
@@ -154,7 +151,7 @@ async def evaluate_text(request: TextEvaluationRequest):
         classifier = TextClassifier()
 
         # Prepare batches
-        batch_size = 24
+        batch_size = 16  # Reduced batch size
         quotes = test_dataset["quote"]
         num_batches = len(quotes) // batch_size + (1 if len(quotes) % batch_size != 0 else 0)
         batches = [
@@ -166,7 +163,7 @@ async def evaluate_text(request: TextEvaluationRequest):
         batch_results = [[] for _ in range(num_batches)]
         
         # Process batches in parallel
-        max_workers = min(os.cpu_count(), 4)
+        max_workers = min(os.cpu_count(), 2)  # Reduced workers
         print(f"Processing with {max_workers} workers")
         
         with ThreadPoolExecutor(max_workers=max_workers) as executor:
@@ -222,5 +219,4 @@ async def evaluate_text(request: TextEvaluationRequest):
 
     except Exception as e:
         print(f"Error in evaluate_text: {str(e)}")
-        raise Exception(f"Failed to process request: {str(e)}")
-    
+        raise Exception(f"Failed to process request: {str(e)}")