import gradio as gr from PyPDF2 import PdfReader from transformers import AutoTokenizer, AutoModel import torch import faiss import numpy as np from groq import Groq import os # ------------- CONSTANTS ------------------------------------------------------ LEGAL_BERT_MODEL = "nlpaueb/legal-bert-base-uncased" # Multiple legal documents - adjust PDFs here DOCS = [ ("bns_full.pdf", "Bharatiya Nyaya Sanhita 2023"), ("bns_ipc_mapping.pdf", "BNS-IPC Comparative Mapping"), ] MAX_CHUNK_SIZE = 1000 OVERLAP = 200 TOP_K = 5 # Number of chunks to retrieve for context LLAMA_MODEL = 'llama-3.3-70b-versatile' # Groq API setup GROQ_API_KEY = os.getenv("GROQ_API_KEY") groq_client = Groq(api_key=GROQ_API_KEY) # ------------- LEGAL-BERT EMBEDDER CLASS ------------------------------------ class LegalBERTEmbedder: def __init__(self, model_name=LEGAL_BERT_MODEL): self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.model = AutoModel.from_pretrained(model_name) self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.model.to(self.device) self.model.eval() def embed(self, texts): all_embeddings = [] with torch.no_grad(): for text in texts: inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(self.device) outputs = self.model(**inputs) cls_embed = outputs.last_hidden_state[:, 0, :].cpu().numpy() all_embeddings.append(cls_embed.flatten()) return np.vstack(all_embeddings) # ------------- PDF PROCESSING FUNCTIONS ------------------------------------ def extract_text_from_pdf(pdf_path): """Extract text from PDF file""" reader = PdfReader(pdf_path) raw_text = "" for page in reader.pages: text = page.extract_text() if text: raw_text += text + "\n" return raw_text def chunk_text(text, max_chunk_size=MAX_CHUNK_SIZE, overlap=OVERLAP): """Split text into overlapping chunks""" chunks = [] start = 0 length = len(text) while start < length: end = min(start + max_chunk_size, length) chunk = text[start:end] chunks.append(chunk) start += max_chunk_size - overlap return chunks # ------------- FAISS INDEX FUNCTIONS --------------------------------------- def build_faiss_index(embeddings): """Build FAISS index for similarity search""" dim = embeddings.shape[1] index = faiss.IndexFlatIP(dim) # Inner product for cosine similarity faiss.normalize_L2(embeddings) index.add(embeddings) return index def query_faiss(index, query_embed, k=TOP_K): """Query FAISS index for top-k similar chunks""" faiss.normalize_L2(query_embed) distances, indices = index.search(query_embed, k) return distances, indices # ------------- LOAD AND PROCESS ALL DOCUMENTS ------------------------------ print("Loading and processing multiple legal documents...") embedder = LegalBERTEmbedder() all_chunks = [] metadata = [] # Store (act_label, original_chunk_text) for reference print("Extracting and chunking text from all PDFs...") for pdf_path, act_label in DOCS: try: raw_text = extract_text_from_pdf(pdf_path) print(f"Extracted {len(raw_text)} characters from {act_label}") chunks = chunk_text(raw_text) print(f"Created {len(chunks)} chunks from {act_label}") # Prefix each chunk with act label for better context labeled_chunks = [f"[{act_label}] {chunk}" for chunk in chunks] all_chunks.extend(labeled_chunks) metadata.extend([(act_label, chunk) for chunk in chunks]) except Exception as e: print(f"Error processing {pdf_path}: {str(e)}") continue print(f"Total chunks created: {len(all_chunks)}") print("Embedding all text chunks with Legal-BERT...") chunk_embeddings = embedder.embed(all_chunks) print("Embeddings created successfully") print("Building FAISS index...") faiss_index = build_faiss_index(chunk_embeddings) print("FAISS index built successfully") # ------------- PROMPT TEMPLATES ------------------------------------------- SYSTEM_PROMPT = """You are a senior Indian legal expert specializing in the Bharatiya Nyaya Sanhita 2023 (BNS) and its correspondence with the Indian Penal Code 1860 (IPC). When answering any question, you MUST use this exact format: CONTEXT/SITUATION: [Provide detailed explanation of the legal context and situation] BNS SECTIONS: [List the specific BNS sections and subsections that apply, with proper citations] IPC SECTIONS (if applicable): [List the corresponding IPC sections based on mappings, with proper citations] SUMMARY: [Provide a clear one-sentence summary highlighting the applicable BNS and IPC sections in **bold** format] Always cite specific sections when available and ensure your response covers relevant BNS provisions and mapped IPC equivalents.""" def build_user_prompt(context, question): """Build the user prompt with context and question""" return f"""Based on the following relevant extracts from BNS and IPC legislation: {context} Question: {question} Please provide a comprehensive legal answer following the exact format specified in the system instructions.""" # ------------- MAIN QUERY FUNCTION ---------------------------------------- def answer_query(user_query): """Main function to answer user queries""" try: # Embed the user query query_embed = embedder.embed([user_query]) # Retrieve top-k similar chunks from FAISS _, indices = query_faiss(faiss_index, query_embed, k=TOP_K) retrieved_chunks = [all_chunks[i] for i in indices[0]] # Prepare context for Llama 3 context = "\n\n".join(retrieved_chunks) # Create chat completion using Groq API with Llama 3 chat_completion = groq_client.chat.completions.create( messages=[ { "role": "system", "content": SYSTEM_PROMPT }, { "role": "user", "content": build_user_prompt(context, user_query) } ], model=LLAMA_MODEL, temperature=0.1, max_tokens=1024 ) return chat_completion.choices[0].message.content.strip() except Exception as e: return f"Error processing query: {str(e)}\n\nPlease check your Groq API key and internet connection." # ------------- GRADIO INTERFACE ------------------------------------------- with gr.Blocks(title="IPC & BNS Legal Assistant") as demo: gr.Markdown(""" # 🏛️ IPC & BNS Legal Assistant **Comprehensive Legal Q&A System covering:** - Bharatiya Nyaya Sanhita 2023 (BNS) - Corresponding Indian Penal Code 1860 (IPC) sections Ask any question about Indian criminal legislation and get structured legal answers with proper citations. """) with gr.Row(): with gr.Column(): query_input = gr.Textbox( label="💼 Enter your legal query", placeholder="e.g., What are the penalties for murder under BNS? What is the IPC equivalent for theft?", lines=4, max_lines=8 ) with gr.Row(): submit_btn = gr.Button("🔍 Get Legal Answer", variant="primary", scale=2) clear_btn = gr.Button("🗑️ Clear", scale=1) with gr.Row(): answer_output = gr.Markdown( label="📋 Legal Analysis", value="*Submit your question to get a structured legal analysis...*" ) # Event handlers submit_btn.click(answer_query, inputs=query_input, outputs=answer_output) query_input.submit(answer_query, inputs=query_input, outputs=answer_output) clear_btn.click(lambda: ("", "*Submit your question to get a structured legal analysis...*"), outputs=[query_input, answer_output]) # Add examples gr.Examples( examples=[ ["What are the penalties for murder under BNS?"], ["What is the IPC equivalent for BNS Section 103?"], ["What constitutes theft according to BNS legislation?"], ["How are punishments defined for assault in BNS?"], ["What are the legal provisions for robbery under IPC and BNS?"] ], inputs=query_input, outputs=answer_output, fn=answer_query, cache_examples=False ) # Launch the interface if __name__ == "__main__": demo.launch( share=False, debug=True, show_error=True )