# src/llm/generator.py

import torch
from .model import model, tokenizer


def generate_answer(prompt, max_new_tokens=50, temperature=0.2):

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            do_sample=False,
            top_p=0.9,
            repetition_penalty=1.2,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.eos_token_id
        )

    response = tokenizer.decode(output[0], skip_special_tokens=True)

    # On enlève le prompt pour ne garder que la réponse
    return response[len(prompt):].strip()