Update app.py

app.py

@@ -2,162 +2,150 @@ import gradio as gr
 from typing import List, Dict
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
-import torch
-from sentence_transformers import SentenceTransformer
+from transformers import pipeline
 import chromadb
 from chromadb.utils import embedding_functions
-import numpy as np
+from sentence_transformers import SentenceTransformer
+import torch
 from tqdm import tqdm
 import os
-from huggingface_hub import login
-from dotenv import load_dotenv
 
-# Load environment variables
-load_dotenv()
-
-# Login to Hugging Face Hub if token is available
-if os.getenv("HUGGINGFACE_API_TOKEN"):
-    login(token=os.getenv("HUGGINGFACE_API_TOKEN"))
-
-class EnhancedChatbot:
+class LegalSearchSystem:
     def __init__(self):
+        print("Initializing Legal Search System...")
+        
         # Initialize ChromaDB
         self.chroma_client = chromadb.Client()
         
-        # Initialize embedding model using sentence-transformers
+        # Initialize embedding function
         self.embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction(
             model_name="all-MiniLM-L6-v2"
         )
         
-        # Create collection with cosine similarity
-        self.collection = self.chroma_client.create_collection(
-            name="text_collection",
-            embedding_function=self.embedding_function,
-            metadata={"hnsw:space": "cosine"}
-        )
-
-        # Initialize the LLM with 8-bit quantization for efficiency
-        model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            load_in_8bit=True,
-            device_map="auto",
-            torch_dtype=torch.float16
-        )
-        
+        # Initialize the model for text generation
         pipe = pipeline(
             "text-generation",
-            model=model,
-            tokenizer=tokenizer,
+            model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
             max_new_tokens=512,
             temperature=0.7,
             top_p=0.95,
-            repetition_penalty=1.15,
-            do_sample=True
+            repetition_penalty=1.15
         )
-        
         self.llm = HuggingFacePipeline(pipeline=pipe)
         
-        # Enhanced prompt templates with specific use cases
+        # Create or get collection
+        self.collection = self.chroma_client.create_collection(
+            name="text_collection",
+            embedding_function=self.embedding_function,
+            metadata={"hnsw:space": "cosine"}
+        )
+        
+        # Initialize chat templates
         self.templates = {
             "default": """
-            You are a knowledgeable assistant providing accurate information based on the given context.
-            
-            GUIDELINES:
-            1. Use ONLY the provided context
-            2. If information is not in context, say "I don't have enough information"
-            3. Be concise and clear
-            4. Use markdown formatting for better readability
-            5. If quoting, use proper citation format
+            You are a legal assistant providing information about the Bharatiya Nyaya Sanhita, 2023.
             
             Context: {context}
             Chat History: {chat_history}
             Question: {question}
             
-            Response:""",
-            
-            "summary": """
-            Create a comprehensive summary of the provided context.
-            
-            Context: {context}
-            
-            REQUIREMENTS:
-            1. Structure the summary with clear headings
-            2. Use bullet points for key information
-            3. Highlight important concepts
-            4. Maintain factual accuracy
-            
-            Summary:""",
-            
-            "technical": """
-            Provide a detailed technical analysis of the context.
-            
-            Context: {context}
-            Question: {question}
-            
-            GUIDELINES:
-            1. Focus on technical specifications
-            2. Explain complex concepts clearly
-            3. Use appropriate technical terminology
-            4. Include relevant examples from context
-            5. Structure the response logically
+            Instructions:
+            1. Answer based ONLY on the provided context
+            2. If information isn't in context, say "I don't have enough information"
+            3. Be precise and cite specific sections when possible
+            4. Use clear, legal terminology
             
-            Technical Analysis:""",
+            Answer:""",
             
-            "comparative": """
-            Compare and analyze different aspects from the context.
+            "summary": """
+            Provide a summary of the legal provisions from the context.
             
             Context: {context}
             Question: {question}
             
-            APPROACH:
-            1. Identify key points for comparison
-            2. Analyze similarities and differences
-            3. Present balanced viewpoints
-            4. Use tables or lists for clarity
+            Format:
+            1. Main Points
+            2. Key Provisions
+            3. Important Definitions
             
-            Comparison:"""
+            Summary:"""
         }
         
         self.chat_history = []
-        self.loaded = False
+        self.initialized = False
 
-    def load_data(self, file_path: str, chunk_size: int = 512, overlap: int = 50):
-        """Load and index data with progress bar"""
-        if self.loaded:
-            return True
-        
+    def initialize_embeddings(self) -> str:
+        """Initialize the system by loading and embedding documents"""
         try:
-            # Read the text file
-            with open(file_path, 'r', encoding='utf-8') as f:
-                content = f.read()
+            if self.initialized:
+                return "System already initialized!"
+            
+            print("Loading documents and creating embeddings...")
+            
+            # Read main text file
+            with open('a2023-45.txt', 'r', encoding='utf-8') as f:
+                text_content = f.read()
+            
+            # Read index file
+            with open('index.txt', 'r', encoding='utf-8') as f:
+                index_lines = f.readlines()
             
-            # Create chunks with overlap
+            # Create chunks
+            chunk_size = 512
             chunks = []
-            for i in range(0, len(content), chunk_size - overlap):
-                chunk = content[i:i + chunk_size]
+            for i in range(0, len(text_content), chunk_size):
+                chunk = text_content[i:i + chunk_size]
                 chunks.append(chunk)
             
-            # Add documents to collection with progress bar
-            for i, chunk in tqdm(enumerate(chunks), desc="Loading chunks", total=len(chunks)):
+            # Add documents to collection
+            print(f"Processing {len(chunks)} chunks...")
+            for i, chunk in enumerate(chunks):
+                # Get corresponding index line if available
+                index_text = index_lines[i].strip() if i < len(index_lines) else f"Chunk {i+1}"
+                
                 self.collection.add(
                     documents=[chunk],
-                    ids=[f"chunk_{i}"],
-                    metadatas=[{"source": file_path, "chunk_id": i}]
+                    ids=[f"doc_{i}"],
+                    metadatas=[{
+                        "index": index_text,
+                        "chunk_number": i
+                    }]
                 )
             
-            self.loaded = True
-            return True
+            self.initialized = True
+            return f"Successfully loaded {len(chunks)} chunks into the system!"
             
         except Exception as e:
-            print(f"Error loading data: {str(e)}")
-            return False
+            return f"Error initializing system: {str(e)}"
 
-    def _search_documents(self, query: str, n_results: int = 5) -> List[Dict]:
+    def verify_system(self) -> str:
+        """Verify system is working properly"""
+        try:
+            # Check document count
+            count = self.collection.count()
+            if count == 0:
+                return "Error: No documents found in the system!"
+            
+            # Test basic query
+            test_query = "What is criminal conspiracy?"
+            results = self.collection.query(
+                query_texts=[test_query],
+                n_results=1
+            )
+            
+            if not results['documents'][0]:
+                return "Error: Search functionality not working properly!"
+            
+            return f"System verification successful! Found {count} documents."
+            
+        except Exception as e:
+            return f"System verification failed: {str(e)}"
+
+    def search(self, query: str, n_results: int = 3) -> List[Dict]:
         """Search for relevant documents"""
+        if not self.initialized:
+            return [{"error": "System not initialized! Please wait."}]
+        
         try:
             results = self.collection.query(
                 query_texts=[query],
@@ -169,7 +157,7 @@ class EnhancedChatbot:
                 {
                     "content": doc,
                     "metadata": meta,
-                    "similarity": 1 - dist  # Convert distance to similarity
+                    "similarity": 1 - dist
                 }
                 for doc, meta, dist in zip(
                     results['documents'][0],
@@ -178,44 +166,30 @@ class EnhancedChatbot:
                 )
             ]
         except Exception as e:
-            print(f"Search error: {str(e)}")
-            return []
-
-    def _select_template(self, query: str) -> str:
-        """Select appropriate template based on query content"""
-        query_lower = query.lower()
-        
-        if any(word in query_lower for word in ["summarize", "summary", "overview"]):
-            return "summary"
-        elif any(word in query_lower for word in ["technical", "explain how", "how does"]):
-            return "technical"
-        elif any(word in query_lower for word in ["compare", "difference", "versus", "vs"]):
-            return "comparative"
-        return "default"
+            return [{"error": f"Search error: {str(e)}"}]
 
     def chat(self, query: str, history) -> str:
-        """Process query and generate response"""
+        """Process query and return response"""
         try:
-            if not self.loaded:
-                if not self.load_data('a2023-45.txt'):
-                    return "Error: Failed to load document data."
+            if not self.initialized:
+                init_msg = self.initialize_embeddings()
+                if "Error" in init_msg:
+                    return init_msg
             
             # Search for relevant content
-            search_results = self._search_documents(query)
+            search_results = self.search(query)
             
-            if not search_results:
-                return "I apologize, but I couldn't find relevant information in the database."
+            if "error" in search_results[0]:
+                return search_results[0]["error"]
             
-            # Prepare context with similarity scores
-            context_parts = []
-            for result in search_results:
-                context_parts.append(
-                    f"[Similarity: {result['similarity']:.2f}]\n{result['content']}"
-                )
-            context = "\n\n".join(context_parts)
+            # Prepare context
+            context = "\n\n".join([
+                f"[Section {r['metadata']['index']}]\n{r['content']}"
+                for r in search_results
+            ])
             
-            # Select and use appropriate template
-            template_type = self._select_template(query)
+            # Select template
+            template_type = "summary" if "summarize" in query.lower() else "default"
             prompt = ChatPromptTemplate.from_template(self.templates[template_type])
             
             # Generate response
@@ -235,38 +209,33 @@ class EnhancedChatbot:
         except Exception as e:
             return f"Error processing query: {str(e)}"
 
-# Initialize chatbot
-chatbot = EnhancedChatbot()
+# Initialize the system
+system = LegalSearchSystem()
 
 # Create Gradio interface
 demo = gr.Interface(
-    fn=chatbot.chat,
+    fn=system.chat,
     inputs=[
         gr.Textbox(
             label="Your Question",
-            placeholder="Ask anything about the document...",
+            placeholder="Ask about the Bharatiya Nyaya Sanhita, 2023...",
             lines=2
         ),
         gr.State([])  # For chat history
     ],
     outputs=gr.Textbox(label="Answer", lines=10),
-    title="🤖 Enhanced Document Q&A System",
+    title="🔍 Bharatiya Nyaya Sanhita, 2023 - Legal Search System",
     description="""
-    ### Advanced Document Question-Answering System
-    
-    **Available Query Types:**
-    - 📝 **General Questions**: Just ask normally
-    - 📊 **Summaries**: Include words like "summarize" or "overview"
-    - 🔧 **Technical Details**: Use words like "technical" or "explain how"
-    - 🔄 **Comparisons**: Ask to "compare" or use "versus"
-    
-    *The system will automatically select the best response format based on your question.*
+    Ask questions about the Bharatiya Nyaya Sanhita, 2023:
+    - For summaries, include the word "summarize" in your question
+    - For specific provisions, ask directly about the topic
+    - System will automatically initialize on first query
     """,
     examples=[
-        ["Can you summarize the main points of the document?"],
-        ["What are the technical details about the implementation?"],
-        ["Compare the different approaches mentioned in the text."],
-        ["What are the key concepts discussed?"]
+        ["What is the definition of criminal conspiracy?"],
+        ["Summarize the provisions related to theft"],
+        ["What are the punishments for corruption?"],
+        ["Explain the concept of culpable homicide"]
     ],
     theme=gr.themes.Soft()
 )