Create llm.py

llm.py

@@ -0,0 +1,10 @@
+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()