Update tasks/text.py

tasks/text.py

@@ -1,6 +1,7 @@
 from fastapi import APIRouter
 from datetime import datetime
 from datasets import load_dataset
+from torch.utils.data import DataLoader
 from sklearn.metrics import accuracy_score
 import random
 
@@ -64,71 +65,82 @@ 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))]
+    # predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
+
+    # Chemins du modèle et du tokenizer
     path_model = 'MatthiasPicard/checkpoint4200_batch16_modern_bert_valloss_0.79_0.74acc'
     path_tokenizer = "answerdotai/ModernBERT-base"
     
-    model = AutoModelForSequenceClassification.from_pretrained(path_model)
+    # Détection du GPU
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Charger le modèle et le tokenizer
+    model = AutoModelForSequenceClassification.from_pretrained(path_model).half().to(device)  # Model en half precision sur GPU
     tokenizer = AutoTokenizer.from_pretrained(path_tokenizer)
     
+    # Fonction de préprocessing
     def preprocess_function(df):
-        return tokenizer(df["quote"], truncation=True)
+        tokenized = tokenizer(df["quote"], truncation=True) # Removed padding here
+        return tokenized
+    
+    # Appliquer le préprocessing
     tokenized_test = test_dataset.map(preprocess_function, batched=True)
-     
-    # training_args = torch.load("training_args.bin")
-    # training_args.eval_strategy='no'
-
-    model = model.half()
-    model.eval() 
     
-    data_collator = DataCollatorWithPadding(tokenizer)
+    # Convertir le dataset au format PyTorch
+    tokenized_test.set_format(type="torch", columns=["input_ids", "attention_mask"])
+    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+    
+    # Créer le DataLoader avec un batch_size > 1 pour optimiser le passage GPU
+    batch_size = 4  # Ajuster selon la mémoire dispo sur GPU
+    test_loader = DataLoader(tokenized_test, batch_size=batch_size, collate_fn=data_collator)
+    
+    model = model.half()
+    model.eval()
     
-    trainer = Trainer(
-        model=model,
-        # args=training_args,
-        tokenizer=tokenizer,
-        data_collator=data_collator
-    )
+    # Inférence sur GPU
+    predictions = []
+    with torch.no_grad():
+        for batch in test_loader:
+            input_ids = batch['input_ids'].to(device)
+            attention_mask = batch['attention_mask'].to(device)
     
-    trainer.args.per_device_eval_batch_size = 32
+            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+            logits = outputs.logits
+            preds = torch.argmax(logits, dim=-1)
     
-    preds = trainer.predict(tokenized_test)
+            predictions.extend(preds.cpu().numpy())  # Remettre sur CPU pour stockage
+
     
-    # path_model = 'MatthiasPi/modernbert_finetunedV1'
+    # path_model = 'MatthiasPicard/checkpoint4200_batch16_modern_bert_valloss_0.79_0.74acc'
     # path_tokenizer = "answerdotai/ModernBERT-base"
     
-    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    # model = AutoModelForSequenceClassification.from_pretrained(path_model).to(device).eval()
+    # model = AutoModelForSequenceClassification.from_pretrained(path_model)
     # tokenizer = AutoTokenizer.from_pretrained(path_tokenizer)
-
-    # model.half()
     
-    # # Use optimized tokenization
     # def preprocess_function(df):
-    #     return tokenizer(df["quote"], truncation=True, padding="max_length")
-
+    #     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'
+
+    # model = model.half()
+    # model.eval() 
+    
+    # data_collator = DataCollatorWithPadding(tokenizer)
+    
+    # trainer = Trainer(
+    #     model=model,
+    #     # args=training_args,
+    #     tokenizer=tokenizer,
+    #     data_collator=data_collator
+    # )
+    
+    # trainer.args.per_device_eval_batch_size = 16
+    
+    # preds = trainer.predict(tokenized_test)
     
-    # # 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, batch_size=16):
-    #     all_preds = []
-    #     with torch.no_grad():  # No gradient computation (saves energy)
-    #         for batch in torch.utils.data.DataLoader(dataset, batch_size=batch_size, 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)
-    # print(predictions)
-    predictions = np.array([np.argmax(x) for x in preds[0]])
+    # predictions = np.array([np.argmax(x) for x in preds[0]])
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE