import torch
from torchmetrics import BLEUScore, METEOR
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

class CodeEvaluator:
    def __init__(self, model_name):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
        self.bleu = BLEUScore()
        self.meteor = METEOR()

    def evaluate(self, nl_input, target_code):
        inputs = self.tokenizer(nl_input, return_tensors="pt").to(self.device)
        outputs = self.model.generate(
            **inputs,
        )
        generated_code = self.tokenizer.decode(outputs[0], skip_special_tokens=True)

        bleu_score = self.bleu(generated_code, target_code)
        meteor_score = self.meteor(generated_code, target_code)
        return bleu_score, meteor_score

if __name__ == "__main__":
    model_name = "S-Dreamer/PyCodeT5"
    evaluator = CodeEvaluator(model_name)

    nl_input = "Write a Python function to calculate the factorial of a number."
    target_code = """
def factorial(n):
  if n == 0:
    return 1
  else:
    return n * factorial(n-1)
"""
    bleu_score, meteor_score = evaluator.evaluate(nl_input, target_code)
    print(f"BLEU score: {bleu_score}")
    print(f"METEOR score: {meteor_score}")