update

tasks/text.py

@@ -5,7 +5,7 @@ from sklearn.metrics import accuracy_score
 import random
 from torch.utils.data import DataLoader
 
-from transformers import AutoTokenizer,BertForSequenceClassification,AutoModelForSequenceClassification,Trainer, TrainingArguments,DataCollatorWithPadding
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments, DataCollatorWithPadding
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
@@ -24,9 +24,8 @@ async def evaluate_text(request: TextEvaluationRequest):
     """
     Evaluate text classification for climate disinformation detection.
     
-    Current Model: Random Baseline
-    - Makes random predictions from the label space (0-7)
-    - Used as a baseline for comparison
+    Current Model: BERT
+    - Uses a pre-trained BERT model for sequence classification
     """
     # Get space info
     username, space_url = get_space_info()
@@ -52,9 +51,7 @@ async def evaluate_text(request: TextEvaluationRequest):
     # Split dataset
     train_test = dataset["train"]
     test_dataset = dataset["test"]
-    print(test_dataset)
-    print(test_dataset['quote'])
-    test_dataset = test_dataset['quote']
+
     # Start tracking emissions
     tracker.start()
     tracker.start_task("inference")
@@ -69,74 +66,48 @@ async def evaluate_text(request: TextEvaluationRequest):
     tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
     model = AutoModelForSequenceClassification.from_pretrained("cococli/bert-base-uncased-frugalai").to(device)
     
-    
-    # def tokenize_frugal(batch, tokenizer):
-    #     return tokenizer(batch, padding=True, truncation=True)
     def tokenize_function(examples):
         return tokenizer(examples["quote"], padding=True, truncation=True, return_tensors='pt')
 
-# Tokenize the test dataset
-    tokenized_test = dataset["test"].map(tokenize_function, batched=True)
-    
-    # tokenized_test = test_dataset.map(lambda batch: tokenize_frugal(batch, tokenizer), batched=True)
-    # tokenized_test = tokenizer(test_dataset, padding = True, truncation= True)
-    
-    # dataloader = DataLoader(tokenized_test, batch_size=16, shuffle=False)
-    # print("Started prediction run")
-
-    # model.eval()
-    # with torch.no_grad():
-    #     predictions = np.array([])
-    #     for batch in dataloader:
-    #         test_input_ids = batch["input_ids"].to(device)
-    #         test_attention_mask = batch["attention_mask"].to(device)
-    #         outputs = model(test_input_ids, test_attention_mask)
-    #         p = torch.argmax(outputs.logits, dim=1)
-    #         predictions = np.append(predictions, p.cpu().numpy())
-    
-    # print("Finished prediction run")
-    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
-    
+    # Tokenize the test dataset
+    tokenized_test = test_dataset.map(tokenize_function, batched=True)
+
     # Create DataLoader
+    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
     dataloader = DataLoader(tokenized_test, batch_size=16, shuffle=False, collate_fn=data_collator)
-    
+
     print("Started prediction run")
-    
+
     # Model inference
     model.eval()
     predictions = np.array([])
-    
+
     with torch.no_grad():
-        
         for batch in dataloader:
             test_input_ids = batch["input_ids"].to(device)
             test_attention_mask = batch["attention_mask"].to(device)
             outputs = model(test_input_ids, test_attention_mask)
             p = torch.argmax(outputs.logits, dim=1)
             predictions = np.append(predictions, p.cpu().numpy())
-       
+
     print("Finished prediction run")
-    
+
     # Ensure "label" column exists in dataset
     print(test_dataset.column_names)  # Debugging step
-    
-    
-    
-    # Make random predictions (placeholder for actual model inference)
+
+    # Extract true labels
     true_labels = test_dataset["label"]
-    # predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE
     #--------------------------------------------------------------------------------------------   
 
-    
     # Stop tracking emissions
     emissions_data = tracker.stop_task()
-    
+
     # Calculate accuracy
     accuracy = accuracy_score(true_labels, predictions)
-    
+
     # Prepare results dictionary
     results = {
         "username": username,
@@ -154,5 +125,5 @@ async def evaluate_text(request: TextEvaluationRequest):
             "test_seed": request.test_seed
         }
     }
-    
-    return results 
+
+    return results