Update app.py

app.py

@@ -1,329 +1,203 @@
 import gradio as gr
-from typing import List, Dict, Tuple
+from typing import List, Dict
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline
+import chromadb
+from chromadb.utils import embedding_functions
 from sentence_transformers import SentenceTransformer
-import torch
 import os
-from astrapy.db import AstraDB
-from dotenv import load_dotenv
-from huggingface_hub import login
-import time
-import logging
-import numpy as np
-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 LegalTextSearchBot:
+class ChromaDBChatbot:
     def __init__(self):
-        try:
-            # Initialize AstraDB connection
-            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(os.getenv("ASTRA_DB_COLLECTION"))
-            
-            # Initialize language model
-            model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-            model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                device_map="auto",
-                torch_dtype=torch.float32,
-            )
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
-            
-            # Initialize text generation pipeline
-            pipe = pipeline(
-                "text-generation",
-                model=model,
-                tokenizer=tokenizer,
-                max_new_tokens=512,
-                temperature=0.7,
-                top_p=0.95,
-                repetition_penalty=1.15,
-                device_map="auto"
-            )
-            self.llm = HuggingFacePipeline(pipeline=pipe)
-            
-            # Initialize sentence transformer for embeddings
-            self.embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
-            
-            self.template = """
-            IMPORTANT: You are a legal assistant that provides accurate information based on the Indian legal sections provided in the context.
+        # Initialize in-memory ChromaDB
+        self.chroma_client = chromadb.Client()
+        
+        # Initialize embedding function
+        self.embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction(
+            model_name="all-MiniLM-L6-v2"
+        )
+        
+        # Create or get collection
+        self.collection = self.chroma_client.create_collection(
+            name="text_collection",
+            embedding_function=self.embedding_function
+        )
+
+        # Initialize the model - using a smaller model suitable for CPU
+        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)
+        
+        # Enhanced prompt templates
+        self.templates = {
+            "default": """
+            IMPORTANT: You are a helpful assistant that provides information based on the retrieved 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."
+            1. Base your response ONLY on the provided context
+            2. If you cannot find relevant information, respond with: "I apologize, but I cannot find information about that in the 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}
+            4. Be concise and accurate in your responses
+            5. If quoting from the context, clearly indicate it
             
+            Context: {context}
             Chat History: {chat_history}
-            
             Question: {question}
             
-            Answer:"""
+            Answer:""",
             
-            self.prompt = ChatPromptTemplate.from_template(self.template)
-            self.chat_history = ""
-            self.is_searching = False
+            "summary": """
+            Create a concise summary of the following context.
             
-            logger.info("Successfully initialized LegalTextSearchBot")
+            Context: {context}
             
-        except Exception as e:
-            logger.error(f"Error initializing LegalTextSearchBot: {str(e)}")
-            raise
-
-    def get_embedding(self, text: str) -> List[float]:
-        """Generate embedding vector for text"""
-        try:
-            # Clean and prepare text
-            text = text.replace('\n', ' ').strip()
-            if not text:
-                text = " "  # Ensure non-empty input
+            Key Requirements:
+            1. Highlight the main points
+            2. Keep it brief and clear
+            3. Use bullet points if appropriate
+            4. Include only information from the context
             
-            # Generate embedding
-            embedding = self.embedding_model.encode(text)
+            Summary:""",
             
-            # Pad or truncate to 1024 dimensions
-            if len(embedding) < 1024:
-                embedding = np.pad(embedding, (0, 1024 - len(embedding)))
-            elif len(embedding) > 1024:
-                embedding = embedding[:1024]
+            "technical": """
+            Provide a technical explanation based on the context.
             
-            return embedding.tolist()
+            Context: {context}
+            Question: {question}
             
-        except Exception as e:
-            logger.error(f"Error generating embedding: {str(e)}")
-            raise
-
-    @lru_cache(maxsize=100)
-    def _cached_search(self, query: str) -> tuple:
-        """Cached version of vector search"""
-        try:
-            # Generate embedding for query
-            query_embedding = self.get_embedding(query)
+            Guidelines:
+            1. Focus on technical details
+            2. Explain complex concepts clearly
+            3. Use appropriate technical terminology
+            4. Provide examples if present in the context
             
-            results = list(self.collection.vector_find(
-                query_embedding,
-                top_k=5,  # Using top_k instead of limit
-                fields=["section_number", "title", "chapter_number", "chapter_title", 
-                       "content", "type", "metadata"]
-            ))
-            return tuple(results)
-        except Exception as e:
-            logger.error(f"Error in vector search: {str(e)}")
-            return tuple()
+            Technical Explanation:"""
+        }
+        
+        self.chat_history = ""
+        self.loaded = False
 
-    def _search_astra(self, query: str) -> List[Dict]:
-        if not self.is_searching:
-            return []
-            
+    def load_data(self, file_path: str):
+        """Load data into ChromaDB"""
+        if self.loaded:
+            return
+        
         try:
-            results = list(self._cached_search(query))
-            
-            if not results and self.is_searching:
-                # Fallback to regular search
-                cursor = self.collection.find({})
-                results = []
-                for doc in cursor:
-                    if len(results) >= 5:
-                        break
-                    results.append(doc)
+            # Read the text file
+            with open(file_path, 'r', encoding='utf-8') as f:
+                content = f.read()
+            
+            # Split into chunks (512 tokens each with 50 token overlap)
+            chunk_size = 512
+            overlap = 50
+            chunks = []
+            
+            for i in range(0, len(content), chunk_size - overlap):
+                chunk = content[i:i + chunk_size]
+                chunks.append(chunk)
+            
+            # Add documents to collection
+            self.collection.add(
+                documents=chunks,
+                ids=[f"doc_{i}" for i in range(len(chunks))]
+            )
             
-            return results
+            self.loaded = True
+            print(f"Loaded {len(chunks)} chunks into ChromaDB")
             
         except Exception as e:
-            logger.error(f"Error searching AstraDB: {str(e)}")
-            return []
+            print(f"Error loading data: {str(e)}")
+            return False
 
-    def format_section(self, section: Dict) -> str:
+    def _search_chroma(self, query: str) -> List[Dict]:
+        """Search ChromaDB for relevant documents"""
         try:
-            return f"""
-{'='*80}
-Chapter {section.get('chapter_number', 'N/A')}: {section.get('chapter_title', 'N/A')}
-Section {section.get('section_number', 'N/A')}: {section.get('title', 'N/A')}
-Type: {section.get('type', 'section')}
-
-Content:
-{section.get('content', 'N/A')}
-
-References: {', '.join(section.get('metadata', {}).get('references', [])) or 'None'}
-{'='*80}
-"""
+            results = self.collection.query(
+                query_texts=[query],
+                n_results=5
+            )
+            return [{"content": doc} for doc in results['documents'][0]]
         except Exception as e:
-            logger.error(f"Error formatting section: {str(e)}")
-            return str(section)
+            print(f"Error searching ChromaDB: {str(e)}")
+            return []
 
-    def generate_ai_response(self, context: str, query: str) -> str:
-        """Generate AI interpretation with error handling"""
+    def chat(self, query: str, history) -> str:
+        """Process a query and return a response"""
         try:
-            chain = self.prompt | self.llm
-            response = chain.invoke({
-                "context": context,
-                "chat_history": self.chat_history,
-                "question": query
-            })
+            if not self.loaded:
+                self.load_data('a2023-45.txt')
             
-            # Handle different response types
-            if isinstance(response, dict):
-                return response.get('text', str(response))
-            elif isinstance(response, list):
-                return response[0] if response else "No response generated"
-            else:
-                return str(response)
-                
-        except Exception as e:
-            logger.error(f"Error generating AI response: {str(e)}")
-            return "I apologize, but I encountered an error while interpreting the legal sections. Please try rephrasing your question."
-
-    def search_sections(self, query: str, progress=gr.Progress()) -> Tuple[str, str]:
-        self.is_searching = True
-        start_time = time.time()
-        
-        try:
-            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."
+            # Determine template type based on query
+            template_type = "default"
+            if any(word in query.lower() for word in ["summarize", "summary"]):
+                template_type = "summary"
+            elif any(word in query.lower() for word in ["technical", "explain", "how does"]):
+                template_type = "technical"
             
-            progress(0.1, desc="Searching relevant sections...")
-            search_results = self._search_astra(query)
+            # Search ChromaDB for relevant content
+            search_results = self._search_chroma(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 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', 'N/A')}: {result.get('title', 'N/A')}
-{result.get('content', 'N/A')}
-""")
-                progress((0.3 + (idx * 0.1)), desc=f"Processing result {idx + 1} of {len(search_results)}...")
+                return "I apologize, but I cannot find information about that in the database."
             
-            if not self.is_searching:
-                return "Search cancelled.", "Search was stopped by user."
+            # Extract and combine relevant content
+            context = "\n\n".join([result['content'] for result in search_results])
             
-            progress(0.8, desc="Generating AI interpretation...")
-            context = "\n\n".join(context_parts)
+            # Create prompt with selected template
+            prompt = ChatPromptTemplate.from_template(self.templates[template_type])
             
-            ai_response = self.generate_ai_response(context, query)
-            self.chat_history += f"\nUser: {query}\nAI: {ai_response}\n"
-            
-            elapsed_time = time.time() - start_time
-            logger.info(f"Search completed in {elapsed_time:.2f} seconds")
+            # Generate response using LLM
+            chain = prompt | self.llm
+            result = chain.invoke({
+                "context": context,
+                "chat_history": self.chat_history,
+                "question": query
+            })
             
-            progress(1.0, desc="Search complete!")
-            return "\n".join(raw_results), ai_response
+            # Update chat history
+            self.chat_history += f"\nUser: {query}\nAI: {result}\n"
             
+            return result
         except Exception as e:
-            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
+            return f"Error processing query: {str(e)}"
 
-    def stop_search(self):
-        """Stop the current search operation"""
-        self.is_searching = False
-        return "Search cancelled.", "Search was stopped by user."
+# Initialize the chatbot
+chatbot = ChromaDBChatbot()
 
-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
-        
-        Search through the Bharatiya Nyaya Sanhita, 2023 and get:
-        1. 📜 Relevant sections, explanations, and illustrations
-        2. 🤖 AI-powered interpretation of the legal content
-        
-        *Use the Stop button if you want to cancel a long-running search.*
-        """)
-        
-        with gr.Row():
-            query_input = gr.Textbox(
-                label="Your Query",
-                placeholder="e.g., What are the penalties for public servants who conceal information?",
-                lines=2
-            )
-        
-        with gr.Row():
-            search_button = gr.Button("🔍 Search", variant="primary", scale=4)
-            stop_button = gr.Button("🛑 Stop", variant="stop", scale=1)
-        
-        with gr.Row():
-            raw_output = gr.Markdown(label="📜 Relevant Legal Sections")
-            ai_output = gr.Markdown(label="🤖 AI Interpretation")
-        
-        gr.Examples(
-            examples=[
-                "What are the penalties for public servants who conceal information?",
-                "What constitutes criminal conspiracy?",
-                "Explain the provisions related to culpable homicide",
-                "What are the penalties for causing death by negligence?",
-                "What are the punishments for corruption?"
-            ],
-            inputs=query_input,
-            label="Example Queries"
-        )
-        
-        # Handle search
-        search_event = search_button.click(
-            fn=search_bot.search_sections,
-            inputs=query_input,
-            outputs=[raw_output, ai_output],
-        )
-        
-        # Handle stop
-        stop_button.click(
-            fn=search_bot.stop_search,
-            outputs=[raw_output, ai_output],
-            cancels=[search_event]
-        )
-        
-        # Handle Enter key
-        query_input.submit(
-            fn=search_bot.search_sections,
-            inputs=query_input,
-            outputs=[raw_output, ai_output],
-        )
-    
-    return iface
+# Create the Gradio interface
+demo = gr.Interface(
+    fn=chatbot.chat,
+    inputs=[
+        gr.Textbox(
+            label="Your Question",
+            placeholder="Ask anything about the document...",
+            lines=2
+        ),
+        gr.State([])  # For chat history
+    ],
+    outputs=gr.Textbox(label="Answer", lines=10),
+    title="ChromaDB-powered Document Q&A",
+    description="""
+    Ask questions about your document:
+    - For summaries, include words like 'summarize' or 'summary'
+    - For technical details, use words like 'technical', 'explain', 'how does'
+    - For general questions, just ask normally
+    """,
+    examples=[
+        ["Can you summarize the main points?"],
+        ["What are the technical details about this topic?"],
+        ["Give me a general overview of the content."],
+    ],
+    theme=gr.themes.Soft()
+)
 
+# Launch the interface
 if __name__ == "__main__":
-    try:
-        demo = create_interface()
-        demo.launch()
-    except Exception as e:
-        logger.error(f"Error launching application: {str(e)}")
-else:
-    try:
-        demo = create_interface()
-        app = demo.launch(share=False)
-    except Exception as e:
-        logger.error(f"Error launching application: {str(e)}")
+    demo.launch()