Update tasks/text.py

tasks/text.py

@@ -2,14 +2,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 DistilBertTokenizerFast, DistilBertForSequenceClassification
+from transformers import Trainer, TrainingArguments
+import torch
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+DESCRIPTION = "DistilBERT Baseline"
 ROUTE = "/text"
 
 @router.post(ROUTE, tags=["Text Task"], 
@@ -18,9 +20,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: DistilBERT
+    - Fine-tunes and evaluates a DistilBERT model on the given dataset
     """
     # Get space info
     username, space_url = get_space_info()
@@ -45,32 +46,62 @@ async def evaluate_text(request: TextEvaluationRequest):
 
     # Split dataset
     train_test = dataset["train"].train_test_split(test_size=request.test_size, seed=request.test_seed)
+    train_dataset = train_test["train"]
     test_dataset = train_test["test"]
-    
+
+    # Tokenizer and model
+    tokenizer = DistilBertTokenizerFast.from_pretrained("distilbert-base-uncased")
+    model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=8)
+
+    # Tokenize datasets
+    def preprocess(examples):
+        return tokenizer(examples["text"], truncation=True, padding=True, max_length=512)
+
+    train_dataset = train_dataset.map(preprocess, batched=True)
+    test_dataset = test_dataset.map(preprocess, batched=True)
+
+    train_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
+    test_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
+
+    # Training arguments
+    training_args = TrainingArguments(
+        output_dir="./results",
+        evaluation_strategy="epoch",
+        learning_rate=5e-5,
+        per_device_train_batch_size=16,
+        per_device_eval_batch_size=16,
+        num_train_epochs=3,
+        weight_decay=0.01,
+        logging_dir="./logs",
+        logging_steps=10,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=test_dataset,
+        tokenizer=tokenizer,
+    )
+
     # 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))]
+    # Train and evaluate the model
+    trainer.train()
 
-    #--------------------------------------------------------------------------------------------
-    # YOUR MODEL INFERENCE STOPS HERE
-    #--------------------------------------------------------------------------------------------   
+    # Perform inference
+    predictions = trainer.predict(test_dataset).predictions
+    predictions = torch.argmax(torch.tensor(predictions), axis=1).tolist()
+    true_labels = test_dataset["label"]
 
-    
     # Stop tracking emissions
     emissions_data = tracker.stop_task()
-    
+
     # Calculate accuracy
     accuracy = accuracy_score(true_labels, predictions)
-    
+
     # Prepare results dictionary
     results = {
         "username": username,
@@ -89,4 +120,4 @@ async def evaluate_text(request: TextEvaluationRequest):
         }
     }
     
-    return results 
+    return results