Update tasks/text.py

tasks/text.py

@@ -16,7 +16,7 @@ DESCRIPTION = "GTE Architecture"
 ROUTE = "/text"
 
 class AutoBertClassifier(nn.Module):
-    def __init__(self, num_labels=8, model_path="Alibaba-NLP/gte-modernbert-base"):
+    def __init__(self, num_labels=8, model_path="TaylorAI/gte-tiny"):
         super().__init__()
         self.tokenizer = AutoTokenizer.from_pretrained(model_path)
         self.bert = AutoModel.from_pretrained(model_path)
@@ -34,7 +34,7 @@ class AutoBertClassifier(nn.Module):
 
 device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
     
-model_repo = "elucidator8918/frugal-ai-text-final"
+model_repo = "elucidator8918/frugal-ai-text-tiny-final"
 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)