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, AutoModelForCausalLM, AutoTokenizer, AutoModel import torch import os from astrapy.db import AstraDB from dotenv import load_dotenv from huggingface_hub import login import time import logging from functools import lru_cache import numpy as np # 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: 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 for text generation 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 embedding model self.embedding_model_name = "sentence-transformers/all-MiniLM-L6-v2" self.embedding_pipeline = pipeline( "feature-extraction", model=self.embedding_model_name, device_map="auto" ) 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 def get_embedding(self, text: str) -> List[float]: """Generate embedding vector for text""" try: # Clean and prepare text text = text.replace('\n', ' ').strip() # Generate embedding outputs = self.embedding_pipeline(text) embeddings = torch.mean(torch.tensor(outputs[0]), dim=0) # Convert to list and ensure correct dimension embedding_list = embeddings.tolist() # Pad or truncate to exactly 1024 dimensions if len(embedding_list) < 1024: embedding_list.extend([0.0] * (1024 - len(embedding_list))) elif len(embedding_list) > 1024: embedding_list = embedding_list[:1024] return embedding_list 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) results = list(self.collection.vector_find( query_embedding, limit=5, 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() def _search_astra(self, query: str) -> List[Dict]: if not self.is_searching: return [] try: results = list(self._cached_search(query)) if not results and self.is_searching: results = list(self.collection.find( {}, limit=5 )) return results except Exception as e: logger.error(f"Error searching AstraDB: {str(e)}") return [] def format_section(self, section: Dict) -> str: 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} """ except Exception as 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.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." 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 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=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) chain = self.prompt | self.llm ai_response = chain.invoke({ "context": context, "chat_history": self.chat_history, "question": 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") progress(1.0, desc="Search complete!") return "\n".join(raw_results), ai_response 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 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 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 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)}")