Update llm.py

llm.py

@@ -1,12 +1,10 @@
-# llm.py
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# load model once
 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()
+    inputs = tokenizer.encode(prompt, return_tensors='pt', truncation=True, max_length=1024)
+    outputs = model.generate(inputs, max_new_tokens=60, do_sample=True, pad_token_id=tokenizer.eos_token_id)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()