Update tasks/text.py

tasks/text.py

@@ -9,7 +9,7 @@ from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+DESCRIPTION = "GTE Architecture"
 ROUTE = "/text"
 
 @router.post(ROUTE, tags=["Text Task"], 
@@ -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: GTE Architecture
     """
     # Get space info
     username, space_url = get_space_info()
@@ -46,6 +44,18 @@ 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"]
+
+    true_labels = test_dataset["label"]
+    texts = test_dataset["quote"]
+
+    model_repo = "elucidator8918/frugal-ai-text"
+    config = AutoConfig.from_pretrained(model_repo)
+    model = AutoModelForSequenceClassification.from_pretrained(model_repo)
+    tokenizer = AutoTokenizer.from_pretrained(model_repo)
+
+    device = torch.device("cuda") if torch.cuda.is_available() torch.device("cpu")
+    model = model.to(device)
+    model.eval()
     
     # Start tracking emissions
     tracker.start()
@@ -56,15 +66,23 @@ async def evaluate_text(request: TextEvaluationRequest):
     # 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))]
+    text_encoding = tokenizer(
+        texts,
+        truncation=True,
+        padding=True,
+        return_tensors="pt",
+    )
+
+    with torch.no_grad():
+        text_input_ids = text_encoding["input_ids"].to(device)
+        text_attention_mask = text_encoding["attention_mask"].to(device)
+        outputs = model(test_input_ids, test_attention_mask)
+        predictions = torch.argmax(outputs.logits, dim=1).cpu().numpy()
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE
     #--------------------------------------------------------------------------------------------   
 
-    
     # Stop tracking emissions
     emissions_data = tracker.stop_task()