llm.py

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")

def generate_answer(context, question):
    prompt = f"Context:\n{context}\n\nQuestion: {question}\nAnswer:"
    inputs = tokenizer.encode(prompt, return_tensors='pt', max_length=1024, truncation=True)
    outputs = model.generate(inputs, max_new_tokens=50, do_sample=True)
    return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()