Update src/generation.py

src/generation.py

@@ -1,29 +1,24 @@
-import numpy as np
-from sentence_transformers import SentenceTransformer
-import json
-import os
-from typing import List, Dict
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+import torch
 
-class DocumentRetriever:
-    def __init__(self, model_name='all-MiniLM-L6-v2'):
-        self.model = SentenceTransformer(model_name)
-        self.documents = self._load_documents()
-        self.doc_embeddings = self._embed_documents()
+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 _load_documents(self) -> List[Dict]:
-        with open('data/rupeia_document.json', 'r') as f:
-            return json.load(f)
+    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 _embed_documents(self) -> np.ndarray:
-        texts = [doc['content'] for doc in self.documents]
-        return self.model.encode(texts)
-
-    def retrieve(self, query: str, top_k: int = 3) -> List[Dict]:
-        query_embedding = self.model.encode(query)
-        scores = np.dot(self.doc_embeddings, query_embedding)
-        top_indices = np.argsort(scores)[-top_k:][::-1]
-        return [self.documents[i] for i in top_indices]
-
-def retrieve_relevant_documents(query: str) -> List[Dict]:
-    retriever = DocumentRetriever()
-    return retriever.retrieve(query)
+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)