Update tasks/text.py

tasks/text.py

@@ -65,27 +65,60 @@ async def evaluate_text(request: TextEvaluationRequest):
     # Make random predictions (placeholder for actual model inference)
     true_labels = test_dataset["label"]
     # predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
+    # path_model = 'MatthiasPi/modernbert_finetunedV1'
+    # path_tokenizer = "answerdotai/ModernBERT-base"
+    
+    # model = AutoModelForSequenceClassification.from_pretrained(path_model)
+    # tokenizer = AutoTokenizer.from_pretrained(path_tokenizer)
+    
+    # def preprocess_function(df):
+    #     return tokenizer(df["quote"], truncation=True)
+    # tokenized_test = test_dataset.map(preprocess_function, batched=True)
+     
+    # # training_args = torch.load("training_args.bin")
+    # # training_args.eval_strategy='no'
+     
+    # trainer = Trainer(
+    #     model=model,
+    #     # args=training_args,
+    #     tokenizer=tokenizer
+    # )
+    
+    # preds = trainer.predict(tokenized_test)
+    
     path_model = 'MatthiasPi/modernbert_finetunedV1'
     path_tokenizer = "answerdotai/ModernBERT-base"
     
-    model = AutoModelForSequenceClassification.from_pretrained(path_model)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = AutoModelForSequenceClassification.from_pretrained(path_model).to(device).eval()
     tokenizer = AutoTokenizer.from_pretrained(path_tokenizer)
     
+    # Use optimized tokenization
     def preprocess_function(df):
-        return tokenizer(df["quote"], truncation=True)
+        return tokenizer(df["quote"], truncation=True, padding="max_length")
+
     tokenized_test = test_dataset.map(preprocess_function, batched=True)
-     
-    # training_args = torch.load("training_args.bin")
-    # training_args.eval_strategy='no'
-     
-    trainer = Trainer(
-        model=model,
-        # args=training_args,
-        tokenizer=tokenizer
-    )
-    
-    preds = trainer.predict(tokenized_test)
-    predictions = np.array([np.argmax(x) for x in preds[0]])
+    
+    # Convert dataset to PyTorch tensors for efficient inference
+    def collate_fn(batch):
+        input_ids = torch.tensor([example["input_ids"] for example in batch]).to(device)
+        attention_mask = torch.tensor([example["attention_mask"] for example in batch]).to(device)
+        return {"input_ids": input_ids, "attention_mask": attention_mask}
+    
+    # Optimized inference function
+    def predict(dataset):
+        all_preds = []
+        with torch.no_grad():  # No gradient computation (saves energy)
+            for batch in torch.utils.data.DataLoader(dataset, batch_size=len(dataset), collate_fn=collate_fn):
+                outputs = model(**batch)
+                preds = torch.argmax(outputs.logits, dim=-1).cpu().numpy()
+                all_preds.extend(preds)
+        return np.array(all_preds)
+    
+    # Run inference
+    predictions = predict(tokenized_test)
+
+    # predictions = np.array([np.argmax(x) for x in preds[0]])
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE