import os import gradio as gr from typing import List, Tuple import numpy as np from pypdf import PdfReader from docx import Document from sentence_transformers import SentenceTransformer import faiss from huggingface_hub import InferenceClient # --- Configuration & Styling --- # SPJIMR Branding Colors PRIMARY_BLUE = "#004890" ACCENT_GOLD = "#C5A059" BG_LIGHT = "#F8F9FA" TEXT_DARK = "#1E1E1E" CUSTOM_CSS = f""" .container {{ max-width: 900px; margin: auto; padding: 20px; font-family: 'Inter', sans-serif; }} .header {{ background-color: {PRIMARY_BLUE}; padding: 20px; border-radius: 12px 12px 0 0; color: white; text-align: center; border-bottom: 4px solid {ACCENT_GOLD}; margin-bottom: 20px; }} .chat-container {{ border-radius: 12px; box-shadow: 0 4px 20px rgba(0,0,0,0.08); background: white; }} .gr-button-primary {{ background-color: {PRIMARY_BLUE} !important; border: none !important; border-radius: 8px !important; color: white !important; }} .gr-button-secondary {{ background-color: white !important; border: 1px solid {PRIMARY_BLUE} !important; border-radius: 8px !important; color: {PRIMARY_BLUE} !important; }} #footer {{ text-align: center; color: #666; margin-top: 20px; font-size: 0.9em; }} .user-bubble {{ background-color: {PRIMARY_BLUE} !important; color: white !important; border-radius: 15px 15px 0 15px !important; }} .bot-bubble {{ background-color: #F1F1F1 !important; color: {TEXT_DARK} !important; border-radius: 15px 15px 15px 0 !important; }} """ # --- Backend Logic --- # Initialize LLM Client (Hugging Face Inference API) # Get token from environment variable 'HF_TOKEN' hf_token = os.getenv("HF_TOKEN") client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3", token=hf_token) # Initialize Embedding Model embed_model = SentenceTransformer("all-MiniLM-L6-v2") class DocumentProcessor: def __init__(self): self.index = None self.chunks = [] self.file_name = "" def extract_text(self, file_path: str) -> str: ext = os.path.splitext(file_path)[1].lower() text = "" if ext == ".pdf": reader = PdfReader(file_path) for page in reader.pages: text += page.extract_text() + "\n" elif ext == ".docx": doc = Document(file_path) for para in doc.paragraphs: text += para.text + "\n" elif ext == ".txt": with open(file_path, "r", encoding="utf-8") as f: text = f.read() return text def chunk_text(self, text: str, chunk_size: int = 500, overlap: int = 50) -> List[str]: words = text.split() chunks = [] for i in range(0, len(words), chunk_size - overlap): chunk = " ".join(words[i : i + chunk_size]) chunks.append(chunk) return chunks def process_document(self, file): if file is None: return "No file uploaded." self.file_name = os.path.basename(file.name) text = self.extract_text(file.name) if not text.strip(): return "Could not extract text from document." self.chunks = self.chunk_text(text) embeddings = embed_model.encode(self.chunks) # Create FAISS Index dimension = embeddings.shape[1] self.index = faiss.IndexFlatL2(dimension) self.index.add(np.array(embeddings).astype("float32")) return f"Successfully processed: {self.file_name} ({len(self.chunks)} chunks)" def retrieve(self, query: str, k: int = 3) -> str: if self.index is None: return "" query_vec = embed_model.encode([query]) distances, indices = self.index.search(np.array(query_vec).astype("float32"), k) context = "" for idx in indices[0]: if idx < len(self.chunks): context += self.chunks[idx] + "\n---\n" return context # Global state for document processing doc_proc = DocumentProcessor() def chat_response(message, history): # System Prompt system_message = "You are a helpful assistant for SPJIMR. " # Check if we have context from a document context = doc_proc.retrieve(message) if context: prompt = f"{system_message}Answer the user question based ONLY on the provided context. If the answer is not in the context, say you don't know based on the document.\n\nContext:\n{context}\n\nQuestion: {message}\nAnswer:" else: prompt = f"{system_message}You are a general assistant. User: {message}\nAssistant:" # Call LLM response = "" try: # Mistral-7B format formatted_prompt = f"[INST] {prompt} [/INST]" for msg in client.text_generation(formatted_prompt, max_new_tokens=512, stream=True): response += msg yield response except Exception as e: yield f"Error connecting to LLM: {str(e)}. Please ensure HF_TOKEN is set." def clear_session(): global doc_proc doc_proc = DocumentProcessor() return None, [], "Session cleared." # --- UI Construction --- with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft(primary_hue="blue")) as demo: with gr.Div(elem_classes="container"): # Header gr.HTML(f"""

Document Assistant

SPJIMR Decision Support System

""") with gr.Row(): # Left Column: Upload & Controls with gr.Column(scale=1): file_input = gr.File( label="Upload Document (PDF, TXT, DOCX)", file_types=[".pdf", ".docx", ".txt"], elem_id="upload-box" ) status_msg = gr.Markdown("Status: Ready", elem_id="status") with gr.Row(): process_btn = gr.Button("Analyze Document", variant="primary") clear_btn = gr.Button("Clear All", variant="secondary") gr.Markdown(""" ### Instructions: 1. Upload a document. 2. Click 'Analyze Document'. 3. Ask questions in the chat! """) # Right Column: Chat Interface with gr.Column(scale=2): chatbot = gr.Chatbot( label="Chat Interface", height=500, bubble_full_width=False, elem_id="chat-window" ) msg_input = gr.Textbox( placeholder="Type your question here...", label=None, container=False, lines=1 ) with gr.Row(): submit_btn = gr.Button("Send", variant="primary") # Footer gr.HTML("") # --- Event Handlers --- def handle_upload(file): if file is None: return "No file selected." msg = doc_proc.process_document(file) return msg process_btn.click( handle_upload, inputs=[file_input], outputs=[status_msg], show_progress=True ) submit_btn.click( chat_response, inputs=[msg_input, chatbot], outputs=[chatbot], show_progress=True ).then( lambda: "", None, [msg_input] ) msg_input.submit( chat_response, inputs=[msg_input, chatbot], outputs=[chatbot], show_progress=True ).then( lambda: "", None, [msg_input] ) clear_btn.click( clear_session, outputs=[file_input, chatbot, status_msg] ) if __name__ == "__main__": demo.launch()