Update tasks/text.py

tasks/text.py

@@ -1,5 +1,6 @@
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
-from fastapi import APIRouter
+from fastapi import FastAPI, APIRouter
+from fastapi.middleware.cors import CORSMiddleware
 from datetime import datetime
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score
@@ -9,12 +10,23 @@ import numpy as np
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
+# Initialize FastAPI app and router
+app = FastAPI()
 router = APIRouter()
 
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
 DESCRIPTION = "Efficient Climate Disinformation Detection"
 ROUTE = "/text"
 
-@router.post(ROUTE, tags=["Text Task"], description=DESCRIPTION)
+@router.post("/text", tags=["Text Task"], description=DESCRIPTION)
 async def evaluate_text(request: TextEvaluationRequest):
     """
     Evaluate text classification for climate disinformation detection.
@@ -46,21 +58,21 @@ async def evaluate_text(request: TextEvaluationRequest):
 
     try:
         # Model configuration
-        model_name = "distilbert-base-uncased"  # Lighter model than MobileBERT
-        BATCH_SIZE = 64  # Increased batch size
-        MAX_LENGTH = 128  # Reduced sequence length
+        model_name = "distilbert-base-uncased"
+        BATCH_SIZE = 64
+        MAX_LENGTH = 128
 
         # Initialize tokenizer and model
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         model = AutoModelForSequenceClassification.from_pretrained(
             model_name,
             num_labels=8,
-            problem_type="single_label_classification",
+            problem_type="single_label_classification"
         )
 
-        # Enable mixed precision training if available
+        # Enable mixed precision if available
         if torch.cuda.is_available():
-            model = model.half()  # Convert to FP16
+            model = model.half()
 
         # Move model to device
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -71,14 +83,14 @@ async def evaluate_text(request: TextEvaluationRequest):
         test_texts = test_dataset["quote"]
         predictions = []
 
-        # Process in efficient batches
+        # Process in batches
         for i in range(0, len(test_texts), BATCH_SIZE):
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
 
             batch_texts = test_texts[i:i + BATCH_SIZE]
             
-            # Efficient tokenization
+            # Tokenize batch
             inputs = tokenizer(
                 batch_texts,
                 padding=True,
@@ -87,18 +99,22 @@ async def evaluate_text(request: TextEvaluationRequest):
                 return_tensors="pt"
             )
             
-            # Move inputs to device efficiently
+            # Move inputs to device
             inputs = {k: v.to(device) for k, v in inputs.items()}
 
-            # Inference with optimizations
+            # Run inference
             with torch.no_grad(), torch.cuda.amp.autocast(enabled=torch.cuda.is_available()):
                 outputs = model(**inputs)
                 batch_preds = torch.argmax(outputs.logits, dim=1)
                 predictions.extend(batch_preds.cpu().numpy())
 
-        # Get true labels and compute accuracy
+        # Get true labels
         true_labels = test_dataset['label']
+
+        # Stop tracking emissions
         emissions_data = tracker.stop_task()
+        
+        # Calculate accuracy
         accuracy = accuracy_score(true_labels, predictions)
         
         # Prepare results
@@ -123,4 +139,12 @@ async def evaluate_text(request: TextEvaluationRequest):
         
     except Exception as e:
         tracker.stop_task()
-        raise e
+        raise e
+
+# Include the router
+app.include_router(router)
+
+# Add a health check endpoint
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}