Update tasks/text.py

tasks/text.py

@@ -12,6 +12,16 @@ from safetensors.torch import load_file
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    
+model_repo = "elucidator8918/frugal-ai-text"
+model = AutoBertClassifier(num_labels=8)
+model.load_state_dict(load_file(hf_hub_download(repo_id=model_repo, filename="model.safetensors")))
+tokenizer = AutoTokenizer.from_pretrained(model_repo)
+
+model = model.to(device)
+model.eval()
+
 router = APIRouter()
 
 DESCRIPTION = "GTE Architecture"
@@ -70,16 +80,6 @@ async def evaluate_text(request: TextEvaluationRequest):
     true_labels = test_dataset["label"]
     texts = test_dataset["quote"]
     
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-    
-    model_repo = "elucidator8918/frugal-ai-text"
-    model = AutoBertClassifier(num_labels=8)
-    model.load_state_dict(load_file(hf_hub_download(repo_id=model_repo, filename="model.safetensors")))
-    tokenizer = AutoTokenizer.from_pretrained(model_repo)
-    
-    model = model.to(device)
-    model.eval()
-    
     # Start tracking emissions
     tracker.start()
     tracker.start_task("inference")
@@ -94,6 +94,7 @@ async def evaluate_text(request: TextEvaluationRequest):
         truncation=True,
         padding=True,
         return_tensors="pt",
+        max_length=256
     )
 
     with torch.no_grad():
@@ -101,7 +102,7 @@ async def evaluate_text(request: TextEvaluationRequest):
         text_attention_mask = text_encoding["attention_mask"].to(device)
         outputs = model(text_input_ids, text_attention_mask)
         predictions = torch.argmax(outputs.logits, dim=1).cpu().numpy()
-
+    
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE
     #--------------------------------------------------------------------------------------------   
@@ -111,6 +112,8 @@ async def evaluate_text(request: TextEvaluationRequest):
     
     # Calculate accuracy
     accuracy = accuracy_score(true_labels, predictions)
+
+    print(f"Accuracy = {accuracy}")
     
     # Prepare results dictionary
     results = {
@@ -129,5 +132,7 @@ async def evaluate_text(request: TextEvaluationRequest):
             "test_seed": request.test_seed
         }
     }
+
+    print(results)
     
     return results