Update tasks/text.py

tasks/text.py

@@ -12,18 +12,6 @@ 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"
 ROUTE = "/text"
 
@@ -44,6 +32,18 @@ class AutoBertClassifier(nn.Module):
         logits = self.classifier(pooled_output)
         return logits
 
+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()
+
 @router.post(ROUTE, tags=["Text Task"], 
              description=DESCRIPTION)
 async def evaluate_text(request: TextEvaluationRequest):