Update tasks/text.py

tasks/text.py

@@ -1,17 +1,16 @@
 from fastapi import APIRouter
 from datetime import datetime
 from datasets import load_dataset
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.naive_bayes import MultinomialNB
 from sklearn.metrics import accuracy_score
-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 = "DistilBERT Baseline"
+DESCRIPTION = "Naive Bayes Text Classifier"
 ROUTE = "/text"
 
 @router.post(ROUTE, tags=["Text Task"], 
@@ -20,8 +19,9 @@ async def evaluate_text(request: TextEvaluationRequest):
     """
     Evaluate text classification for climate disinformation detection.
     
-    Current Model: DistilBERT
-    - Fine-tunes and evaluates a DistilBERT model on the given dataset
+    Current Model: Naive Bayes Classifier
+    - Uses TF-IDF for text vectorization
+    - Trains and evaluates a Multinomial Naive Bayes model
     """
     # Get space info
     username, space_url = get_space_info()
@@ -49,58 +49,33 @@ async def evaluate_text(request: TextEvaluationRequest):
     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,
-    )
+    # Extract text and labels
+    train_texts = [x["text"] for x in train_dataset]
+    train_labels = [x["label"] for x in train_dataset]
+    test_texts = [x["text"] for x in test_dataset]
+    test_labels = [x["label"] for x in test_dataset]
+
+    # TF-IDF Vectorization
+    vectorizer = TfidfVectorizer(max_features=5000)
+    train_vectors = vectorizer.fit_transform(train_texts)
+    test_vectors = vectorizer.transform(test_texts)
+
+    # Train Naive Bayes Classifier
+    model = MultinomialNB()
+    model.fit(train_vectors, train_labels)
 
     # Start tracking emissions
     tracker.start()
     tracker.start_task("inference")
 
-    # Train and evaluate the model
-    trainer.train()
-
-    # Perform inference
-    predictions = trainer.predict(test_dataset).predictions
-    predictions = torch.argmax(torch.tensor(predictions), axis=1).tolist()
-    true_labels = test_dataset["label"]
+    # Inference
+    predictions = model.predict(test_vectors)
 
     # Stop tracking emissions
     emissions_data = tracker.stop_task()
 
     # Calculate accuracy
-    accuracy = accuracy_score(true_labels, predictions)
+    accuracy = accuracy_score(test_labels, predictions)
 
     # Prepare results dictionary
     results = {