Update tasks/text.py

tasks/text.py

@@ -9,7 +9,13 @@ from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+import numpy as np
+
+from transformers import AutoTokenizer,BertForSequenceClassification,AutoModelForSequenceClassification,Trainer, TrainingArguments,DataCollatorWithPadding
+from datasets import Dataset
+import torch
+
+DESCRIPTION = "BERT-Fine Tune"
 ROUTE = "/text"
 
 @router.post(ROUTE, tags=["Text Task"], 
@@ -56,10 +62,27 @@ async def evaluate_text(request: TextEvaluationRequest):
     # Update the code below to replace the random baseline by your model inference within the inference pass where the energy consumption and emissions are tracked.
     #--------------------------------------------------------------------------------------------   
     
-    # Make random predictions (placeholder for actual model inference)
-    true_labels = test_dataset["label"]
-    predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
+    model = BertForSequenceClassification.from_pretrained("Oriaz/climate_change_bert_classif")
+    tokenizer = AutoTokenizer.from_pretrained("Oriaz/climate_change_bert_classif")
+
+    ## 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]])
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE
     #--------------------------------------------------------------------------------------------