Update app/query_rewriter.py

app/query_rewriter.py

@@ -1,9 +1,9 @@
 import os
-import ollama
 import logging
 import sys
+from transformers import pipeline
 
-# Configure logging for stdout only
+# Configure logging
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
@@ -13,50 +13,75 @@ logger = logging.getLogger(__name__)
 
 class QueryRewriter:
     def __init__(self):
-        self.model = os.getenv('OLLAMA_MODEL', "phi3")
-        self.ollama_host = os.getenv('OLLAMA_HOST', 'http://ollama:11434')
+        self.model = pipeline(
+            "text-generation",
+            model="google/flan-t5-base",  # Using a smaller model suitable for Spaces
+            device=-1  # Use CPU
+        )
+        logger.info("Initialized QueryRewriter with flan-t5-base model")
 
     def generate(self, prompt):
         try:
-            response = ollama.chat(
-                model=self.model,
-                messages=[{"role": "user", "content": prompt}]
-            )
-            return response['message']['content']
+            response = self.model(
+                prompt,
+                max_length=256,
+                min_length=32,
+                num_return_sequences=1
+            )[0]['generated_text']
+            return response
         except Exception as e:
             logger.error(f"Error generating response: {e}")
-            return f"Error: {str(e)}"
+            return None
 
     def rewrite_cot(self, query):
         prompt = f"""
-        Rewrite the following query using Chain-of-Thought reasoning:
-        Query: {query}
+        Rewrite the following query using step-by-step reasoning:
+        
+        Original query: {query}
+        
+        Steps:
+        1. What is the main question being asked?
+        2. What are the key components?
+        3. How can we make it clearer?
         
         Rewritten query:
         """
+        
         rewritten_query = self.generate(prompt)
-        if rewritten_query.startswith("Error:"):
-            logger.error(f"Error in CoT rewriting: {rewritten_query}")
+        if rewritten_query is None:
+            logger.error(f"Error in CoT rewriting for query: {query}")
             return query, prompt  # Return original query if rewriting fails
-        return rewritten_query, prompt
+            
+        # Extract the rewritten query (everything after "Rewritten query:")
+        try:
+            final_query = rewritten_query.split("Rewritten query:")[-1].strip()
+            return final_query, prompt
+        except Exception as e:
+            logger.error(f"Error extracting rewritten query: {e}")
+            return query, prompt
 
     def rewrite_react(self, query):
         prompt = f"""
-        Rewrite the following query using the ReAct framework (Reasoning and Acting):
-        Query: {query}
+        Rewrite the following query using a systematic approach:
         
-        Thought 1:
-        Action 1:
-        Observation 1:
+        Original query: {query}
         
-        Thought 2:
-        Action 2:
-        Observation 2:
+        Thought: What information are we looking for?
+        Action: Break down the query into key components
+        Observation: Identify the main focus
         
-        Final rewritten query:
+        Rewritten query:
         """
+        
         rewritten_query = self.generate(prompt)
-        if rewritten_query.startswith("Error:"):
-            logger.error(f"Error in ReAct rewriting: {rewritten_query}")
-            return query, prompt  # Return original query if rewriting fails
-        return rewritten_query, prompt
+        if rewritten_query is None:
+            logger.error(f"Error in ReAct rewriting for query: {query}")
+            return query, prompt
+            
+        # Extract the rewritten query
+        try:
+            final_query = rewritten_query.split("Rewritten query:")[-1].strip()
+            return final_query, prompt
+        except Exception as e:
+            logger.error(f"Error extracting rewritten query: {e}")
+            return query, prompt