Update tasks/text.py

tasks/text.py

@@ -1,16 +1,16 @@
-from fastapi import APIRouter
-from datetime import datetime
-from datasets import load_dataset
-from sklearn.metrics import accuracy_score
-import random
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+import os
+import zipfile
 
-from .utils.evaluation import TextEvaluationRequest
-from .utils.emissions import tracker, clean_emissions_data, get_space_info
+# Unzip the uploaded file
+model_dir = "/tasks/bert_fine_tuned_model"
+with zipfile.ZipFile('/tasks/bert_fine_tuned_model-20250107T090607Z-001.zip', 'r') as zip_ref:
+    zip_ref.extractall(model_dir)
 
-router = APIRouter()
-
-DESCRIPTION = "Random Baseline"
-ROUTE = "/text"
+# Load the fine-tuned BERT model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_dir)
+model = AutoModelForSequenceClassification.from_pretrained(model_dir)
 
 @router.post(ROUTE, tags=["Text Task"], 
              description=DESCRIPTION)
@@ -18,9 +18,7 @@ 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: Fine-Tuned BERT
     """
     # Get space info
     username, space_url = get_space_info()
@@ -46,31 +44,37 @@ async def evaluate_text(request: TextEvaluationRequest):
     # Split dataset
     train_test = dataset["train"].train_test_split(test_size=request.test_size, seed=request.test_seed)
     test_dataset = train_test["test"]
+
+    # Preprocess the test dataset
+    def preprocess_function(examples):
+        return tokenizer(examples["text"], truncation=True, padding=True, max_length=512)
     
+    test_dataset = test_dataset.map(preprocess_function, batched=True)
+
+    # Convert to PyTorch dataset
+    test_dataset = test_dataset.with_format("torch")
+
+    # Assign device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+
     # Start tracking emissions
     tracker.start()
     tracker.start_task("inference")
 
-    #--------------------------------------------------------------------------------------------
-    # YOUR MODEL INFERENCE CODE HERE
-    # Update the code below to replace the random baseline by your model inference within the inference pass where the energy consumption and emissions are tracked.
-    #--------------------------------------------------------------------------------------------   
-    
-    # Make random predictions (placeholder for actual model inference)
-    true_labels = test_dataset["label"]
-    predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
-
-    #--------------------------------------------------------------------------------------------
-    # YOUR MODEL INFERENCE STOPS HERE
-    #--------------------------------------------------------------------------------------------   
+    # Perform inference
+    true_labels = test_dataset["label"].tolist()
+    inputs = {key: test_dataset[key].to(device) for key in ["input_ids", "attention_mask"]}
+    with torch.no_grad():
+        logits = model(**inputs).logits
+    predictions = torch.argmax(logits, dim=1).tolist()
 
-    
     # Stop tracking emissions
     emissions_data = tracker.stop_task()
-    
+
     # Calculate accuracy
     accuracy = accuracy_score(true_labels, predictions)
-    
+
     # Prepare results dictionary
     results = {
         "username": username,
@@ -89,4 +93,4 @@ async def evaluate_text(request: TextEvaluationRequest):
         }
     }
     
-    return results 
+    return results