Update src/generation.py

src/generation.py

@@ -1,24 +1,10 @@
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
-import torch
+from transformers import pipeline
 
 class ResponseGenerator:
     def __init__(self, model_name='distilgpt2'):
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModelForCausalLM.from_pretrained(model_name)
-        self.tokenizer.pad_token = self.tokenizer.eos_token
-
-    def generate(self, prompt: str, max_length: int = 150) -> str:
-        inputs = self.tokenizer(prompt, return_tensors="pt")
-        outputs = self.model.generate(
-            inputs.input_ids,
-            max_length=max_length,
-            num_return_sequences=1,
-            pad_token_id=self.tokenizer.eos_token_id
-        )
-        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-def generate_response(query: str, context: list) -> str:
-    generator = ResponseGenerator()
-    context_str = "\n".join([doc['content'] for doc in context])
-    prompt = f"Question: {query}\nContext: {context_str}\nAnswer:"
-    return generator.generate(prompt)
+        self.generator = pipeline('text-generation', model=model_name)
+    
+    def generate(self, prompt: str, context: list, max_length: int = 150):
+        context_str = "\n".join([doc['content'] for doc in context])
+        full_prompt = f"Context: {context_str}\nQuestion: {prompt}\nAnswer:"
+        return self.generator(full_prompt, max_length=max_length, num_return_sequences=1)[0]['generated_text']