import wandb
import yaml
from transformers import pipeline

# Charger la configuration
with open('config/config.yaml', 'r') as f:
    config = yaml.safe_load(f)

# Initialiser wandb
wandb.init(project=config['wandb']['project'], entity=config['wandb']['entity'])

# Charger le modèle fine-tuned
model_name = "results_student"  # Remplacer par le chemin vers le modèle student
tokenizer_name = "distilbert-base-uncased"

# Configuration du pipeline
nlp = pipeline("text-classification", model=model_name, tokenizer=tokenizer_name)

# Simuler des exemples pour l'évaluation
examples = [
    {"reference": "This is a great movie.", "candidate": "This is a fantastic movie."},
    {"reference": "I love this film.", "candidate": "I enjoy this movie."}
]

def evaluate_prompt(example, shots=0):
    prompt = example["candidate"]
    if shots == 1:
        prompt = "Classify the sentiment of the following text: " + prompt
    elif shots > 1:
        prompt = "Classify the sentiment of the following text based on these examples:\n" + \
                 "Example: This movie is terrible. -> Negative\n" + \
                 "Example: I love this movie. -> Positive\n" + \
                 prompt

    result = nlp(prompt)[0]
    return result

# Évaluer les prompts
for example in examples:
    for shots in [0, 1, 5]:
        result = evaluate_prompt(example, shots)
        wandb.log({
            'example': example['candidate'],
            'shots': shots,
            'result': result
        })