Update app.py

app.py

@@ -164,37 +164,39 @@ class RAGPipeline:
         query = re.sub(r'\s+', ' ', query)
         return query
 
-    def postprocess_response(self, response: str) -> str:
-        """Clean up the generated response"""
-        response = response.strip()
-        response = re.sub(r'\s+', ' ', response)
-        response = re.sub(r'\d{4}-\d{2}-\d{2}\s\d{2}:\d{2}:\d{2}(?:\+\d{2}:?\d{2})?', '', response)
-        return response
+    ### Added on Nov 2, 2024
 
-    def query_model(self, prompt: str) -> str:
-        """Query the local Llama model"""
-        try:
-            if self.llm is None:
-                raise RuntimeError("Model not initialized")
+    # def postprocess_response(self, response: str) -> str:
+    #     """Clean up the generated response"""
+    #     response = response.strip()
+    #     response = re.sub(r'\s+', ' ', response)
+    #     response = re.sub(r'\d{4}-\d{2}-\d{2}\s\d{2}:\d{2}:\d{2}(?:\+\d{2}:?\d{2})?', '', response)
+    #     return response
+
+    # def query_model(self, prompt: str) -> str:
+    #     """Query the local Llama model"""
+    #     try:
+    #         if self.llm is None:
+    #             raise RuntimeError("Model not initialized")
                 
-            response = self.llm(
-                prompt,
-                max_tokens=512,
-                temperature=0.4,
-                top_p=0.95,
-                echo=False,
-                stop=["Question:", "\n\n"]
-            )
-            
-            if response and 'choices' in response and len(response['choices']) > 0:
-                text = response['choices'][0].get('text', '').strip()
-                return text
-            else:
-                raise ValueError("No valid response generated")
+    #         response = self.llm(
+    #             prompt,
+    #             max_tokens=512,
+    #             temperature=0.4,
+    #             top_p=0.95,
+    #             echo=False,
+    #             stop=["Question:", "\n\n"]
+    #         )
+            
+    #         if response and 'choices' in response and len(response['choices']) > 0:
+    #             text = response['choices'][0].get('text', '').strip()
+    #             return text
+    #         else:
+    #             raise ValueError("No valid response generated")
                 
-        except Exception as e:
-            logging.error(f"Error in query_model: {str(e)}")
-            raise
+    #     except Exception as e:
+    #         logging.error(f"Error in query_model: {str(e)}")
+    #         raise
 
     # def process_query(self, query: str, placeholder) -> str:
     #     try:
@@ -258,90 +260,213 @@ class RAGPipeline:
     #         placeholder.warning(message)
     #         return message
 
+    # def process_query(self, query: str, placeholder) -> str:
+    #     try:
+    #         # Preprocess query
+    #         query = self.preprocess_query(query)
+    #         logging.info(f"Processing query: {query}")
+        
+    #         # Show retrieval status
+    #         status = placeholder.empty()
+    #         status.write("🔍 Finding relevant information...")
+        
+    #         # Get embeddings and search
+    #         query_embedding = self.retriever.encode([query])
+    #         similarities = F.cosine_similarity(query_embedding, self.retriever.doc_embeddings)
+    #         scores, indices = torch.topk(similarities, k=min(self.k, len(self.documents)))
+        
+    #         # Log similarity scores
+    #         for idx, score in zip(indices.tolist(), scores.tolist()):
+    #             logging.info(f"Score: {score:.4f} | Document: {self.documents[idx][:100]}...")
+        
+    #         relevant_docs = [self.documents[idx] for idx in indices.tolist()]
+        
+    #         # Update status
+    #         status.write("💭 Generating response...")
+        
+    #         # Prepare context and prompt
+    #         context = "\n".join(relevant_docs[:3])
+    #         prompt = f"""Context information is below:
+    #         {context}
+        
+    #         Given the context above, please answer the following question:
+    #         {query}
+
+    #         Guidelines:
+    #         - If you cannot answer based on the context, say so politely
+    #         - Keep the response concise and focused
+    #         - Only include sports-related information
+    #         - No dates or timestamps in the response
+    #         - Use clear, natural language
+        
+    #         Answer:"""
+        
+    #         # Generate response
+    #         response_placeholder = placeholder.empty()
+        
+    #         try:
+    #             # Add logging for model state
+    #             logging.info("Model state check - Is None?: " + str(self.llm is None))
+            
+    #             # Directly use Llama model
+    #             response = self.llm(
+    #                 prompt,
+    #                 max_tokens=512,
+    #                 temperature=0.4,
+    #                 top_p=0.95,
+    #                 echo=False,
+    #                 stop=["Question:", "\n\n"]
+    #             )
+            
+    #             logging.info(f"Raw model response: {response}")
+            
+    #             if response and isinstance(response, dict) and 'choices' in response:
+    #                 generated_text = response['choices'][0].get('text', '').strip()
+    #                 if generated_text:
+    #                     final_response = self.postprocess_response(generated_text)
+    #                     response_placeholder.markdown(final_response)
+    #                     return final_response
+            
+    #             message = "No relevant answer found. Please try rephrasing your question."
+    #             response_placeholder.warning(message)
+    #             return message
+            
+    #         except Exception as e:
+    #             logging.error(f"Generation error: {str(e)}")
+    #             logging.error(f"Full error details: ", exc_info=True)
+    #             message = f"Had some trouble generating the response: {str(e)}"
+    #             response_placeholder.warning(message)
+    #             return message
+            
+    #     except Exception as e:
+    #         logging.error(f"Process error: {str(e)}")
+    #         logging.error(f"Full error details: ", exc_info=True)
+    #         message = f"Something went wrong: {str(e)}"
+    #         placeholder.warning(message)
+    #         return message
+
+    ### Added on Nov 2, 2024
+    def postprocess_response(self, response: str) -> str:
+    """Clean up the generated response"""
+    try:
+        # Remove datetime patterns and other unwanted content
+        response = re.sub(r'\d{4}-\d{2}-\d{2}(?:T|\s)\d{2}:\d{2}:\d{2}(?:\.\d+)?(?:Z|[+-]\d{2}:?\d{2})?', '', response)
+        response = re.sub(r'User \d+:.*?(?=User \d+:|$)', '', response)
+        response = re.sub(r'\d{2}:\d{2}(?::\d{2})?(?:\s?(?:AM|PM))?', '', response)
+        response = re.sub(r'\d{1,2}[-/]\d{1,2}[-/]\d{2,4}', '', response)
+        response = re.sub(r'(?m)^User \d+:', '', response)
+        
+        # Clean up spacing but preserve intentional paragraph breaks
+        # Replace multiple newlines with two newlines (one paragraph break)
+        response = re.sub(r'\n\s*\n\s*\n+', '\n\n', response)
+        # Replace multiple spaces with single space
+        response = re.sub(r' +', ' ', response)
+        # Clean up beginning/end
+        response = response.strip()
+        
+        return response
+    except Exception as e:
+        logging.error(f"Error in postprocess_response: {str(e)}")
+        return response
+
     def process_query(self, query: str, placeholder) -> str:
         try:
-            # Preprocess query
             query = self.preprocess_query(query)
-            logging.info(f"Processing query: {query}")
-        
-            # Show retrieval status
             status = placeholder.empty()
             status.write("🔍 Finding relevant information...")
-        
-            # Get embeddings and search
+            
             query_embedding = self.retriever.encode([query])
             similarities = F.cosine_similarity(query_embedding, self.retriever.doc_embeddings)
             scores, indices = torch.topk(similarities, k=min(self.k, len(self.documents)))
-        
-            # Log similarity scores
-            for idx, score in zip(indices.tolist(), scores.tolist()):
-                logging.info(f"Score: {score:.4f} | Document: {self.documents[idx][:100]}...")
-        
+            
             relevant_docs = [self.documents[idx] for idx in indices.tolist()]
+            cleaned_docs = []
+            for doc in relevant_docs[:3]:
+                cleaned_text = self.postprocess_response(doc)
+                if cleaned_text:
+                    cleaned_docs.append(cleaned_text)
         
-            # Update status
             status.write("💭 Generating response...")
-        
-            # Prepare context and prompt
-            context = "\n".join(relevant_docs[:3])
+            
             prompt = f"""Context information is below:
-            {context}
-        
+            {' '.join(cleaned_docs)}
+            
             Given the context above, please answer the following question:
             {query}
-
-            Guidelines:
+        
+            Guidelines for your response:
+            - Structure your response in clear, logical paragraphs
+            - Start a new paragraph for each new main point or aspect
+            - If listing multiple items, use separate paragraphs
+            - Keep each paragraph focused on a single topic or point
+            - Use natural paragraph breaks where the content shifts focus
+            - Maintain clear transitions between paragraphs
+            - If providing statistics or achievements, group them logically
+            - If describing different aspects (e.g., career, playing style, achievements), use separate paragraphs
+            - Keep paragraphs concise but complete
+            - Exclude any dates, timestamps, or user comments
+            - Focus on factual sports information
             - If you cannot answer based on the context, say so politely
-            - Keep the response concise and focused
-            - Only include sports-related information
-            - No dates or timestamps in the response
-            - Use clear, natural language
         
+            Format your response with proper paragraph breaks where appropriate.
+            
             Answer:"""
-        
-            # Generate response
+            
             response_placeholder = placeholder.empty()
-        
-            try:
-                # Add logging for model state
-                logging.info("Model state check - Is None?: " + str(self.llm is None))
-            
-                # Directly use Llama model
-                response = self.llm(
-                    prompt,
-                    max_tokens=512,
-                    temperature=0.4,
-                    top_p=0.95,
-                    echo=False,
-                    stop=["Question:", "\n\n"]
-                )
-            
-                logging.info(f"Raw model response: {response}")
-            
-                if response and isinstance(response, dict) and 'choices' in response:
-                    generated_text = response['choices'][0].get('text', '').strip()
-                    if generated_text:
-                        final_response = self.postprocess_response(generated_text)
-                        response_placeholder.markdown(final_response)
-                        return final_response
-            
-                message = "No relevant answer found. Please try rephrasing your question."
-                response_placeholder.warning(message)
-                return message
             
+            try:
+                response_text = self.query_model(prompt)
+                if response_text:
+                    # Clean up the response while preserving paragraph structure
+                    final_response = self.postprocess_response(response_text)
+                    
+                    # Convert cleaned response to markdown with proper paragraph spacing
+                    markdown_response = final_response.replace('\n\n', '\n\n \n\n')  # Add visual spacing between paragraphs
+                    
+                    response_placeholder.markdown(markdown_response)
+                    return final_response
+                else:
+                    message = "No relevant answer found. Please try rephrasing your question."
+                    response_placeholder.warning(message)
+                    return message
+                    
             except Exception as e:
                 logging.error(f"Generation error: {str(e)}")
-                logging.error(f"Full error details: ", exc_info=True)
-                message = f"Had some trouble generating the response: {str(e)}"
+                message = "Had some trouble generating the response. Please try again."
                 response_placeholder.warning(message)
                 return message
-            
+                
         except Exception as e:
             logging.error(f"Process error: {str(e)}")
-            logging.error(f"Full error details: ", exc_info=True)
-            message = f"Something went wrong: {str(e)}"
+            message = "Something went wrong. Please try again with a different question."
             placeholder.warning(message)
-            return message
+            return messag
+
+    def query_model(self, prompt: str) -> str:
+        """Query the local Llama model"""
+        try:
+            if self.llm is None:
+                raise RuntimeError("Model not initialized")
+                
+            response = self.llm(
+                prompt,
+                max_tokens=512,
+                temperature=0.4,
+                top_p=0.95,
+                echo=False,
+                stop=["Question:", "Context:", "Guidelines:"],  # Removed \n\n from stop tokens to allow paragraphs
+                repeat_penalty=1.1  # Added to encourage more diverse text
+            )
+            
+            if response and 'choices' in response and len(response['choices']) > 0:
+                text = response['choices'][0].get('text', '').strip()
+                return text
+            else:
+                raise ValueError("No valid response generated")
+                
+        except Exception as e:
+            logging.error(f"Error in query_model: {str(e)}")
+            raise
 
 @st.cache_resource(show_spinner=False)
 def initialize_rag_pipeline():