Create modules/qa_module.py

modules/qa_module.py

@@ -1,35 +1,59 @@
+
+
 # modules/qa_module.py
+from transformers import pipeline
 from typing import Dict, List
-import chroma
-from langchain import OpenAI
-from core.base_module import AIModule
+import torch
 
-class QAModule(AIModule):
-    def __init__(self, model_name: str = "gpt-3.5-turbo"):
-        self.model = OpenAI(model_name=model_name)
-        self.vector_store = chromadb.Client()
-        self.collection = self.vector_store.create_collection("qa_collection")
-    
-    async def process(self, input_data: Dict) -> Dict:
-        query = input_data.get("query")
-        results = self.collection.query(
-            query_texts=[query],
-            n_results=2
+class EnhancedQAModule:
+    def __init__(
+        self, 
+        model_name: str = "HuggingFaceH4/zephyr-7b-beta",
+        device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    ):
+        self.model = pipeline(
+            "question-answering",
+            model=model_name,
+            device=device,
+            model_kwargs={"torch_dtype": torch.float16 if device == "cuda" else torch.float32}
+        )
+        
+        self.prompt_template = """
+        <|system|>
+        Answer the question based on the provided context. Be concise and specific.
+        If the answer cannot be found in the context, say so.
+        </s>
+        <|user|>
+        Context:
+        {context}
+        
+        Question: {question}
+        </s>
+        <|assistant|>
+        """
+
+    async def process(self, query: str, context_docs: List[Dict]) -> Dict:
+        # Combine context documents
+        context = "\n".join([f"[{doc['metadata']['source']}]: {doc['content']}" 
+                           for doc in context_docs])
+        
+        # Format prompt
+        prompt = self.prompt_template.format(
+            context=context,
+            question=query
         )
-        context = results['documents'][0]
         
-        response = self.model.predict(
-            f"Context: {context}\nQuestion: {query}\nAnswer:"
+        # Generate answer
+        response = self.model(
+            question=query,
+            context=context,
+            max_length=200,
+            num_beams=4,
+            temperature=0.7
         )
         
         return {
-            "answer": response,
-            "sources": context
-        }
-    
-    async def get_status(self) -> Dict:
-        return {"status": "operational", "documents_indexed": len(self.collection)}
-    
-    @property
-    def capabilities(self) -> List[str]:
-        return ["question-answering", "context-aware-responses"]
+            "answer": response["answer"],
+            "confidence": response["score"],
+            "sources": [doc["metadata"]["source"] for doc in context_docs]
+        }