add bert model code

tasks/text.py

@@ -4,8 +4,11 @@ from datasets import load_dataset
 from sklearn.metrics import accuracy_score
 import random
 from skops.io import load
-# Textpreprocessor defined in this scope
-
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoConfig
+import torch
+from torch.utils.data import DataLoader, Dataset
+import numpy as np
+from accelerate.test_utils.testing import get_backend
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
@@ -19,11 +22,10 @@ ROUTE = "/text"
 models_descriptions = {
     "baseline": "random baseline",
     "tfidf_xgb": "TF-IDF vectorizer and XGBoost classifier",
+    "bert_base_pruned": "Pruned BERT base model",
 }
 
 
-# Some code borrowed from Nonnormalizable
-
 def baseline_model(dataset_length: int):
     # Make random predictions (placeholder for actual model inference)
     predictions = [random.randint(0, 7) for _ in range(dataset_length)]
@@ -48,10 +50,59 @@ def tree_classifier(test_dataset: dict, model: str):
 
     return predictions
 
+class TextDataset(Dataset):
+    def __init__(self, texts, tokenizer, max_length=256):
+        self.texts = texts
+        self.tokenized_texts = tokenizer(
+            texts,
+            truncation=True,
+            padding=True,
+            max_length=max_length,
+            return_tensors="pt",
+        )
+
+    def __getitem__(self, idx):
+        item = {key: val[idx] for key, val in self.tokenized_texts.items()}
+        return item
+
+    def __len__(self) -> int:
+        return len(self.texts)
+        
+
+
+def bert_classifier(test_dataset: dict, model: str):
+    texts = test_dataset["quote"]
+
+    model_repo = f"theterryzhang/frugal_ai_{model}"
+
+    model = AutoModelForSequenceClassification.from_pretrained(model_repo)
+    tokenizer = AutoTokenizer.from_pretrained(model_repo)
+
+    # Use CUDA if available 
+    device, _, _ = get_backend()
+
+    model = model.to(device)
+
+    # Prepare dataset
+    dataset = TextDataset(texts, tokenizer=tokenizer)
+    dataloader = DataLoader(dataset, batch_size=32, shuffle=False)
+
+    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())
+    
+    return predictions
+
 
 @router.post(ROUTE, tags=["Text Task"])
 async def evaluate_text(request: TextEvaluationRequest,
-                        model: str = "tfidf_xgb"):
+                        model: str = "bert_base_pruned"):
     """
     Evaluate text classification for climate disinformation detection.
     
@@ -100,6 +151,8 @@ async def evaluate_text(request: TextEvaluationRequest,
         predictions = baseline_model(len(true_labels))
     elif model == "tfidf_xgb":
         predictions = tree_classifier(test_dataset, model='xgb_pipeline')
+    elif 'bert' in model:
+        predictions = bert_classifier(test_dataset, model)
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE