import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained('./')
model = AutoModelForCausalLM.from_pretrained('./')

def test_model(input_text, expected_output):
    inputs = tokenizer(input_text, return_tensors='pt')
    outputs = model.generate(**inputs, max_new_tokens=50)
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    assert generated_text == expected_output, f"Expected: {expected_output}, but got: {generated_text}"

test_model("Salut comment vas tu?", "Salut comment vas tu")