Update llm.py

llm.py

@@ -1,10 +1,23 @@
+# llm.py
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
 model = AutoModelForCausalLM.from_pretrained("distilgpt2")
 
+# Fix: add pad_token_id if missing
+if tokenizer.pad_token_id is None:
+    tokenizer.pad_token_id = tokenizer.eos_token_id
+
 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)
+    # Limit to last N chars if prompt is too long
+    prompt = prompt[-1000:]  
+
+    inputs = tokenizer(prompt, return_tensors='pt', truncation=True, max_length=1024)
+    outputs = model.generate(
+        inputs["input_ids"],
+        max_new_tokens=50,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id  # fix warning
+    )
     return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()