Update tasks/text.py

tasks/text.py

@@ -3,15 +3,20 @@ from datetime import datetime
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score
 import random
-import torch
+
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
-from transformers import AutoTokenizer,BertForSequenceClassification,AutoModelForSequenceClassification,Trainer, TrainingArguments,DataCollatorWithPadding
+
+from peft import PeftModel
+from transformers import AutoTokenizer,AutoModelForSequenceClassification,Trainer, TrainingArguments,DataCollatorWithPadding, BitsAndBytesConfig
+from datasets import Dataset
+import torch
 import numpy as np
 
+
 router = APIRouter()
 
-DESCRIPTION = "ModernBert Baseline"
+DESCRIPTION = "qwen_finetuned"
 ROUTE = "/text"
 
 @router.post(ROUTE, tags=["Text Task"], 
@@ -46,8 +51,8 @@ async def evaluate_text(request: TextEvaluationRequest):
     dataset = dataset.map(lambda x: {"label": LABEL_MAPPING[x["label"]]})
 
     # Split dataset
-    train_test = dataset["train"].train_test_split(test_size=request.test_size, seed=request.test_seed)
-    test_dataset = train_test["test"]
+    train_test = dataset["train"]
+    test_dataset = dataset["test"]
     
     # Start tracking emissions
     tracker.start()
@@ -61,32 +66,50 @@ 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_adapter = 'MatthiasPicard/Qwen3B_model_test'
+    path_model = "Qwen/Qwen2.5-3B-Instruct"
 
-    #--------------------------------------------------------------------------------------------
-    # YOUR MODEL INFERENCE STOPS HERE
-    #--------------------------------------------------------------------------------------------   
-    ## Model loading
-    model = AutoModelForSequenceClassification.from_pretrained("Rcarvalo/test_modernbert_finetuned_v2")
-    tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")
+    bnb_config = BitsAndBytesConfig(
+    load_in_8bit=True
+    )
+    
+    base_model = AutoModelForSequenceClassification.from_pretrained(
+    path_model,
+    num_labels=len(LABEL_MAPPING),
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+    quantization_config=bnb_config
+    )    
 
-    ## Data prep
+    model = PeftModel.from_pretrained(base_model, path_adapter)
+    model.eval()
+    tokenizer = AutoTokenizer.from_pretrained(path_model)
+    
     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'
-
+     
+    # training_args = torch.load("training_args.bin")
+    # training_args.eval_strategy='no'
+     
     trainer = Trainer(
         model=model,
-        args=training_args,
         tokenizer=tokenizer
     )
-
-    ## prediction
+    
     preds = trainer.predict(tokenized_test)
+    
+    
+    
+    # Run inference
+    # predictions = predict(tokenized_test)
+    # print(predictions)
     predictions = np.array([np.argmax(x) for x in preds[0]])
+
+    #--------------------------------------------------------------------------------------------
+    # YOUR MODEL INFERENCE STOPS HERE
+    #--------------------------------------------------------------------------------------------   
+
     
     # Stop tracking emissions
     emissions_data = tracker.stop_task()