Update app.py

app.py

@@ -2,68 +2,90 @@ import gradio as gr
 from typing import List, Dict, Tuple
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
-from transformers import pipeline
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+import torch
 import os
 from astrapy.db import AstraDB
 from dotenv import load_dotenv
 from huggingface_hub import login
 import time
-import threading
-from queue import Queue
-import asyncio
+import logging
+from functools import lru_cache
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
 
 # Load environment variables
 load_dotenv()
 login(token=os.getenv("HUGGINGFACE_API_TOKEN"))
 
-class SearchCancelled(Exception):
-    pass
+# Initialize model with optimized settings
+model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16,
+    device_map="auto",
+    load_in_8bit=True
+)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 class LegalTextSearchBot:
     def __init__(self):
-        self.astra_db = AstraDB(
-            token=os.getenv("ASTRA_DB_APPLICATION_TOKEN"),
-            api_endpoint=os.getenv("ASTRA_DB_API_ENDPOINT")
-        )
-        self.collection = self.astra_db.collection("legal_content")
-        
-        pipe = pipeline(
-            "text-generation",
-            model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
-            max_new_tokens=512,
-            temperature=0.7,
-            top_p=0.95,
-            repetition_penalty=1.15
-        )
-        self.llm = HuggingFacePipeline(pipeline=pipe)
-        
-        self.template = """
-        IMPORTANT: You are a legal assistant that provides accurate information based on the Indian legal sections provided in the context.
-        
-        STRICT RULES:
-        1. Base your response ONLY on the provided legal sections
-        2. If you cannot find relevant information, respond with: "I apologize, but I cannot find information about that in the legal database."
-        3. Do not make assumptions or use external knowledge
-        4. Always cite the specific section numbers you're referring to
-        5. Be precise and accurate in your legal interpretations
-        6. If quoting from the sections, use quotes and cite the section number
-        
-        Context (Legal Sections): {context}
-        
-        Chat History: {chat_history}
-        
-        Question: {question}
-        
-        Answer:"""
-        
-        self.prompt = ChatPromptTemplate.from_template(self.template)
-        self.chat_history = ""
-        self.cancel_search = False
-
-    def _search_astra(self, query: str) -> List[Dict]:
-        if self.cancel_search:
-            raise SearchCancelled("Search was cancelled by user")
+        try:
+            self.astra_db = AstraDB(
+                token=os.getenv("ASTRA_DB_APPLICATION_TOKEN"),
+                api_endpoint=os.getenv("ASTRA_DB_API_ENDPOINT")
+            )
+            self.collection = self.astra_db.collection("legal_content")
+            
+            # Initialize pipeline with optimized settings
+            pipe = pipeline(
+                "text-generation",
+                model=model,
+                tokenizer=tokenizer,
+                max_new_tokens=512,
+                temperature=0.7,
+                top_p=0.95,
+                repetition_penalty=1.15,
+                torch_dtype=torch.float16,
+                device_map="auto"
+            )
+            self.llm = HuggingFacePipeline(pipeline=pipe)
+            
+            self.template = """
+            IMPORTANT: You are a legal assistant that provides accurate information based on the Indian legal sections provided in the context.
+            
+            STRICT RULES:
+            1. Base your response ONLY on the provided legal sections
+            2. If you cannot find relevant information, respond with: "I apologize, but I cannot find information about that in the legal database."
+            3. Do not make assumptions or use external knowledge
+            4. Always cite the specific section numbers you're referring to
+            5. Be precise and accurate in your legal interpretations
+            6. If quoting from the sections, use quotes and cite the section number
+            
+            Context (Legal Sections): {context}
+            
+            Chat History: {chat_history}
+            
+            Question: {question}
             
+            Answer:"""
+            
+            self.prompt = ChatPromptTemplate.from_template(self.template)
+            self.chat_history = ""
+            self.is_searching = False
+            
+        except Exception as e:
+            logger.error(f"Error initializing LegalTextSearchBot: {str(e)}")
+            raise
+
+    @lru_cache(maxsize=100)
+    def _cached_search(self, query: str) -> tuple:
+        """Cached version of vector search to improve performance"""
         try:
             results = list(self.collection.vector_find(
                 query,
@@ -71,8 +93,19 @@ class LegalTextSearchBot:
                 fields=["section_number", "title", "chapter_number", "chapter_title", 
                        "content", "type", "metadata"]
             ))
+            return tuple(results)  # Convert to tuple for caching
+        except Exception as e:
+            logger.error(f"Error in vector search: {str(e)}")
+            return tuple()
+
+    def _search_astra(self, query: str) -> List[Dict]:
+        if not self.is_searching:
+            return []
             
-            if not results and not self.cancel_search:
+        try:
+            results = list(self._cached_search(query))
+            
+            if not results and self.is_searching:
                 results = list(self.collection.find(
                     {},
                     limit=5
@@ -81,14 +114,10 @@ class LegalTextSearchBot:
             return results
             
         except Exception as e:
-            if not isinstance(e, SearchCancelled):
-                print(f"Error searching AstraDB: {str(e)}")
-            raise
+            logger.error(f"Error searching AstraDB: {str(e)}")
+            return []
 
     def format_section(self, section: Dict) -> str:
-        if self.cancel_search:
-            raise SearchCancelled("Search was cancelled by user")
-            
         try:
             return f"""
 {'='*80}
@@ -103,33 +132,44 @@ References: {', '.join(section.get('metadata', {}).get('references', [])) or 'No
 {'='*80}
 """
         except Exception as e:
-            print(f"Error formatting section: {str(e)}")
+            logger.error(f"Error formatting section: {str(e)}")
             return str(section)
 
     def search_sections(self, query: str, progress=gr.Progress()) -> Tuple[str, str]:
-        self.cancel_search = False
+        self.is_searching = True
+        start_time = time.time()
         
         try:
-            progress(0, desc="Searching relevant sections...")
+            progress(0, desc="Initializing search...")
+            if not query.strip():
+                return "Please enter a search query.", "Please provide a specific legal question or topic to search for."
+            
+            progress(0.1, desc="Searching relevant sections...")
             search_results = self._search_astra(query)
             
             if not search_results:
                 return "No relevant sections found.", "I apologize, but I cannot find relevant sections in the database."
             
+            if not self.is_searching:
+                return "Search cancelled.", "Search was stopped by user."
+            
             progress(0.3, desc="Processing results...")
             raw_results = []
             context_parts = []
             
             for idx, result in enumerate(search_results):
-                if self.cancel_search:
-                    raise SearchCancelled("Search was cancelled by user")
+                if not self.is_searching:
+                    return "Search cancelled.", "Search was stopped by user."
                     
                 raw_results.append(self.format_section(result))
                 context_parts.append(f"""
 Section {result.get('section_number')}: {result.get('title')}
 {result.get('content', '')}
 """)
-                progress((0.3 + (idx * 0.1)), desc="Processing results...")
+                progress((0.3 + (idx * 0.1)), desc=f"Processing result {idx + 1} of {len(search_results)}...")
+            
+            if not self.is_searching:
+                return "Search cancelled.", "Search was stopped by user."
             
             progress(0.8, desc="Generating AI interpretation...")
             context = "\n\n".join(context_parts)
@@ -143,21 +183,27 @@ Section {result.get('section_number')}: {result.get('title')}
             
             self.chat_history += f"\nUser: {query}\nAI: {ai_response}\n"
             
-            progress(1.0, desc="Complete!")
+            elapsed_time = time.time() - start_time
+            logger.info(f"Search completed in {elapsed_time:.2f} seconds")
+            
+            progress(1.0, desc="Search complete!")
             return "\n".join(raw_results), ai_response
             
-        except SearchCancelled:
-            return "Search cancelled by user.", "Search was stopped. Please try again with a new query."
         except Exception as e:
-            error_msg = f"Error processing query: {str(e)}"
-            print(error_msg)
-            return error_msg, "An error occurred while processing your query."
+            logger.error(f"Error processing query: {str(e)}")
+            return f"Error processing query: {str(e)}", "An error occurred while processing your query."
+        finally:
+            self.is_searching = False
 
-    def cancel(self):
-        self.cancel_search = True
+    def stop_search(self):
+        """Stop the current search operation"""
+        self.is_searching = False
+        return "Search cancelled.", "Search was stopped by user."
 
 def create_interface():
     with gr.Blocks(title="Bharatiya Nyaya Sanhita Search", theme=gr.themes.Soft()) as iface:
+        search_bot = LegalTextSearchBot()
+        
         gr.Markdown("""
         # 📚 Bharatiya Nyaya Sanhita Legal Search System
         
@@ -165,11 +211,9 @@ def create_interface():
         1. 📜 Relevant sections, explanations, and illustrations
         2. 🤖 AI-powered interpretation of the legal content
         
-        Enter your legal query below:
+        *Use the Stop button if you want to cancel a long-running search.*
         """)
         
-        search_bot = LegalTextSearchBot()
-        
         with gr.Row():
             query_input = gr.Textbox(
                 label="Your Query",
@@ -178,20 +222,12 @@ def create_interface():
             )
         
         with gr.Row():
-            with gr.Column(scale=4):
-                search_button = gr.Button("🔍 Search Legal Sections", variant="primary")
-            with gr.Column(scale=1):
-                stop_button = gr.Button("🛑 Stop Search", variant="stop")
+            search_button = gr.Button("🔍 Search", variant="primary", scale=4)
+            stop_button = gr.Button("🛑 Stop", variant="stop", scale=1)
         
         with gr.Row():
-            with gr.Column():
-                raw_output = gr.Markdown(
-                    label="📜 Relevant Legal Sections"
-                )
-            with gr.Column():
-                ai_output = gr.Markdown(
-                    label="🤖 AI Interpretation"
-                )
+            raw_output = gr.Markdown(label="📜 Relevant Legal Sections")
+            ai_output = gr.Markdown(label="🤖 AI Interpretation")
         
         gr.Examples(
             examples=[
@@ -205,38 +241,38 @@ def create_interface():
             label="Example Queries"
         )
         
-        def search(query):
-            return search_bot.search_sections(query)
-        
-        def stop_search():
-            search_bot.cancel()
-            return "Search cancelled.", "Search stopped by user."
-        
-        search_button.click(
-            fn=search,
+        # Handle search
+        search_event = search_button.click(
+            fn=search_bot.search_sections,
             inputs=query_input,
             outputs=[raw_output, ai_output],
-            cancels=[stop_button]  # Cancel any ongoing search when stop is clicked
         )
         
+        # Handle stop
         stop_button.click(
-            fn=stop_search,
+            fn=search_bot.stop_search,
             outputs=[raw_output, ai_output],
-            cancels=[search_button]  # Cancel the search button when stop is clicked
+            cancels=[search_event]
         )
         
+        # Handle Enter key
         query_input.submit(
-            fn=search,
+            fn=search_bot.search_sections,
             inputs=query_input,
             outputs=[raw_output, ai_output],
-            cancels=[stop_button]
         )
     
     return iface
 
 if __name__ == "__main__":
-    demo = create_interface()
-    demo.launch()
+    try:
+        demo = create_interface()
+        demo.launch()
+    except Exception as e:
+        logger.error(f"Error launching application: {str(e)}")
 else:
-    demo = create_interface()
-    app = demo.launch(share=False)
+    try:
+        demo = create_interface()
+        app = demo.launch(share=False)
+    except Exception as e:
+        logger.error(f"Error launching application: {str(e)}")