require

scripts/evaluate.py

@@ -1,24 +1,28 @@
-from transformers import pipeline, AutoModelForSequenceClassification
+import torch
+from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score, f1_score
 
-# Load dataset
-dataset = load_dataset("allocine")["test"]
+# Load dataset (limit to 10 samples for faster evaluation)
+dataset = load_dataset("allocine")["test"].select(range(10))
+
+# Use GPU if available
+device = 0 if torch.cuda.is_available() else -1
 
 # Load model and tokenizer
 model_path = "./models"
-classifier = pipeline("text-classification", model=model_path, tokenizer=model_path)
+classifier = pipeline("text-classification", model=model_path, tokenizer=model_path, device=device)
 
-# Get actual model labels
+# Load model to get dynamic label mapping
 model = AutoModelForSequenceClassification.from_pretrained(model_path)
-label_map = {v: k for k, v in model.config.label2id.items()}  # Adjust dynamically
+label_map = {v: f"LABEL_{k}" for k, v in model.config.label2id.items()}  # Ensure mapping is correct
 
 # Get predictions
 predictions = [classifier(text["review"], truncation=True, max_length=512)[0]["label"] for text in dataset]
 labels = dataset["label"]
 
 # Convert labels
-predictions = [label_map[p] for p in predictions]
+predictions = [int(label_map[p].split("_")[-1]) for p in predictions]  # Convert back to int labels
 
 # Compute metrics
 accuracy = accuracy_score(labels, predictions)