modernBert

tasks/text.py

@@ -6,10 +6,11 @@ import random
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
+from transformers import AutoModelForSequenceClassification,AutoTokenizer
 
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+DESCRIPTION = "ModernBert Baseline"
 ROUTE = "/text"
 
 @router.post(ROUTE, tags=["Text Task"], 
@@ -63,7 +64,28 @@ async def evaluate_text(request: TextEvaluationRequest):
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE
     #--------------------------------------------------------------------------------------------   
+    ## Model loading
+    model = AutoModelForSequenceClassification.from_pretrained("Rcarvalo/test_modernbert_finetuned_v2")
+    tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")
 
+    ## Data prep
+    def preprocess_function(df):
+        return tokenizer(df["quote"], truncation=True)
+    tokenized_test = test_dataset.map(preprocess_function, batched=True)
+
+    ## Modify inference model
+    training_args = torch.load("./tasks/utils/training_args.bin")
+    training_args.eval_strategy='no'
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        tokenizer=tokenizer
+    )
+
+    ## prediction
+    preds = trainer.predict(tokenized_test)
+    predictions = np.array([np.argmax(x) for x in preds[0]])
     
     # Stop tracking emissions
     emissions_data = tracker.stop_task()