""" TEJASWI.AI — Women's Legal Rights Assistant (India) RAG pipeline: FAISS retrieval over legal PDFs + Qwen2.5-3B-Instruct fine-tuned with a LoRA adapter, served via Gradio on a CPU Space. """ import os import pickle import re from dataclasses import dataclass import spaces import faiss import numpy as np import torch import gradio as gr from sentence_transformers import SentenceTransformer from transformers import AutoTokenizer, AutoModelForCausalLM from peft import PeftModel from huggingface_hub import hf_hub_download # ------------------------------------------------------------------ # Config # ------------------------------------------------------------------ BASE_MODEL_ID = "Qwen/Qwen2.5-3B-Instruct" ADAPTER_ID = "jpo89/qwen2.5-3b-math-dapo" # your fine-tuned LoRA adapter KB_REPO_ID = "jpo89/TEJASWI.AI" # this Space's own repo (KB lives in /kb) CREATOR_NAME = "Jyotirmoy Bhunia" # ------------------------------------------------------------------ # Data structures (must match what was pickled in chunks.pkl) # ------------------------------------------------------------------ @dataclass class DocChunk: text: str source: str page: int # ------------------------------------------------------------------ # Load knowledge base (FAISS index + chunks) from this Space's own repo # ------------------------------------------------------------------ print("Downloading knowledge base files...") index_path = hf_hub_download(repo_id=KB_REPO_ID, filename="kb/index.faiss", repo_type="space") chunks_path = hf_hub_download(repo_id=KB_REPO_ID, filename="kb/chunks.pkl", repo_type="space") index = faiss.read_index(index_path) with open(chunks_path, "rb") as f: chunks = pickle.load(f) print(f"Loaded FAISS index with {index.ntotal} vectors and {len(chunks)} chunks.") # ------------------------------------------------------------------ # Load hero background illustration from this Space's own repo # (uploaded separately to assets/hero_bg.png) # ------------------------------------------------------------------ import base64 HERO_BG_DATA_URI = "" try: hero_img_path = hf_hub_download(repo_id=KB_REPO_ID, filename="assets/hero_bg.png", repo_type="space") with open(hero_img_path, "rb") as image_file: encoded_string = base64.b64encode(image_file.read()).decode("utf-8") HERO_BG_DATA_URI = f"data:image/png;base64,{encoded_string}" print("Hero background image loaded.") except Exception as e: print(f"Hero background image not found ({e}) — continuing without it.") # ------------------------------------------------------------------ # Load embedding model (multilingual e5, CPU) # ------------------------------------------------------------------ print("Loading embedding model...") embedder = SentenceTransformer("intfloat/multilingual-e5-base", device="cpu") def embed_texts(texts, prefix="passage: "): prefixed = [prefix + t for t in texts] return embedder.encode(prefixed, batch_size=32, show_progress_bar=False, normalize_embeddings=True) def retrieve(query: str, k: int = 5): q_emb = embed_texts([query], prefix="query: ") q_emb = np.array(q_emb, dtype="float32") scores, idxs = index.search(q_emb, k) results = [] for score, i in zip(scores[0], idxs[0]): if i == -1: continue c = chunks[i] results.append({"text": c.text, "source": c.source, "page": c.page, "score": float(score)}) return results # ------------------------------------------------------------------ # Model + adapter are NOT loaded at module scope. ZeroGPU has no GPU # attached until a @spaces.GPU function actually runs, and even # CPU-targeted loads inside peft/safetensors can trip CUDA checks # internally. So the entire load happens lazily, inside ask(), on # the first request, and is cached afterward. # ------------------------------------------------------------------ print("Loading tokenizer (lightweight, safe at module scope)...") tokenizer = AutoTokenizer.from_pretrained(ADAPTER_ID) _model = None def get_model(): global _model if _model is None: print("First request: loading base model + LoRA adapter onto GPU...") base = AutoModelForCausalLM.from_pretrained( BASE_MODEL_ID, torch_dtype=torch.bfloat16, ).to("cuda") _model = PeftModel.from_pretrained(base, ADAPTER_ID) _model.eval() print("Model loaded and ready on GPU.") return _model # ------------------------------------------------------------------ # Prompt template # ------------------------------------------------------------------ PROMPT_TEMPLATE = """You are a compassionate, knowledgeable legal information assistant focused on Women's Legal Rights in India. Guidelines: 1. Speak with empathy and respect. Acknowledge the person's situation without being dramatic or presumptuous about facts they haven't shared. 2. Base your answer strictly on the Context below. Do not invent Acts, Sections, case names, or statistics that are not present in the Context. 3. When citing legal provisions (Acts, Sections, Articles, case law), name them clearly and then explain what they mean in plain language. 4. Never suggest that a person's account of harassment, violence, or abuse might be exaggerated, fabricated, or "misused" — your role is to inform her of her rights and options, not to evaluate the credibility of her situation. 5. If the Context does not contain enough information to fully answer the question, say so plainly, share what partial information IS available, and recommend she consult a qualified lawyer, the National Commission for Women, or a local women's helpline for case-specific advice. 6. Keep the response focused and practical: what the law says, what it means for her, and a clear next step. Context: {context} Question: {question} Response:""" def build_context(retrieved_chunks): parts = [] for r in retrieved_chunks: parts.append(f"[Source: {r['source']}, page {r['page']}]\n{r['text']}") return "\n\n---\n\n".join(parts) @spaces.GPU(duration=120) def ask(question: str, k: int = 5, max_new_tokens: int = 384): model = get_model() retrieved = retrieve(question, k=k) context = build_context(retrieved) prompt = PROMPT_TEMPLATE.format(context=context, question=question) messages = [{"role": "user", "content": prompt}] encoded = tokenizer.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt", return_dict=True, ) encoded = {k_: v.to(model.device) for k_, v in encoded.items()} input_ids = encoded["input_ids"] with torch.no_grad(): output = model.generate( **encoded, max_new_tokens=max_new_tokens, do_sample=True, temperature=0.4, top_p=0.9, repetition_penalty=1.1, pad_token_id=tokenizer.eos_token_id, ) generated = output[0][input_ids.shape[-1]:] answer = tokenizer.decode(generated, skip_special_tokens=True).strip() return {"answer": answer, "sources": retrieved} # ------------------------------------------------------------------ # Gradio UI # ------------------------------------------------------------------ EXAMPLE_QUESTIONS = [ "What protection does the Domestic Violence Act 2005 provide?", "What does Section 498-A of the IPC cover?", "What constitutional articles protect women's rights in India?", "Is there a law against workplace harassment for women?", ] CUSTOM_CSS = f""" @import url('https://fonts.googleapis.com/css2?family=Lora:wght@500;600;700&family=Source+Sans+3:wght@400;500;600&family=JetBrains+Mono:wght@400;500&display=swap'); :root {{ --ink-teal: #1B3A3F; --sand: #F7F0E3; --charcoal: #2B2B2B; --sage: #8FA68E; --gold: #D4A857; --hairline: rgba(27, 58, 63, 0.15); --purple-theme: #8E44AD; }} .gradio-container {{ background: #ffffff !important; font-family: 'Source Sans 3', sans-serif !important; max-width: 880px !important; margin: 0 auto !important; }} #hero {{ text-align: left; padding: 56px 24px 56px 32px; position: relative; background-color: #ffffff; background-image: url('{HERO_BG_DATA_URI}'); background-repeat: no-repeat; background-position: right center; background-size: contain; border-radius: 18px; border: 1px solid var(--hairline); overflow: hidden; min-height: 280px; }} #hero h1 {{ font-family: 'Lora', serif !important; font-weight: 700 !important; font-size: 2.6rem !important; color: var(--ink-teal) !important; margin-bottom: 10px !important; letter-spacing: -0.01em; max-width: 55%; text-shadow: 0 0 15px rgba(255, 255, 255, 1), 0 0 30px rgba(255, 255, 255, 0.9); }} #hero p.tagline {{ font-size: 1.05rem; color: var(--charcoal); opacity: 0.95; max-width: 55%; line-height: 1.5; text-shadow: 0 0 15px rgba(255, 255, 255, 1), 0 0 30px rgba(255, 255, 255, 0.9); }} @media (max-width: 768px) {{ #hero {{ background-position: center top; background-size: cover; text-align: center; padding: 30px 15px; }} #hero h1, #hero p.tagline {{ max-width: 100%; background: rgba(255, 255, 255, 0.85); border-radius: 8px; padding: 5px; }} }} .book-divider {{ width: 100%; height: 1px; background: var(--hairline); margin: 32px 0 28px 0; position: relative; }} .book-divider::before {{ content: ""; position: absolute; left: 50%; top: -3px; width: 7px; height: 7px; border-radius: 50%; background: var(--gold); transform: translateX(-50%); }} #ask-panel {{ background: #FFFFFF; border: 1px solid var(--hairline); border-radius: 14px; padding: 28px; box-shadow: 0 2px 18px rgba(27, 58, 63, 0.06); }} #ask-panel label {{ font-family: 'Lora', serif !important; font-weight: 600 !important; color: var(--ink-teal) !important; font-size: 1.1rem !important; }} textarea, input[type="text"] {{ font-family: 'Source Sans 3', sans-serif !important; font-size: 1rem !important; border-radius: 10px !important; border: 1.5px solid var(--hairline) !important; background: #FAFAFA !important; color: var(--charcoal) !important; }} textarea:focus, input[type="text"]:focus {{ border-color: var(--purple-theme) !important; box-shadow: 0 0 0 3px rgba(142, 68, 173, 0.15) !important; outline: none !important; }} button.primary, #ask-button {{ background: var(--purple-theme) !important; color: #FFFFFF !important; border: none !important; border-radius: 10px !important; font-weight: 600 !important; font-family: 'Source Sans 3', sans-serif !important; padding: 12px 28px !important; font-size: 1rem !important; transition: transform 0.15s ease, box-shadow 0.15s ease; }} button.primary:hover, #ask-button:hover {{ transform: translateY(-1px); box-shadow: 0 4px 14px rgba(142, 68, 173, 0.35); }} #answer-box {{ background: #FFFFFF; border: 1px solid var(--hairline); border-radius: 14px; padding: 28px; margin-top: 20px; font-size: 1.02rem; line-height: 1.7; color: var(--charcoal) !important; }} #answer-box * {{ color: var(--charcoal) !important; }} #answer-box h1, #answer-box h2, #answer-box h3 {{ font-family: 'Lora', serif !important; color: var(--ink-teal) !important; }} #sources-accordion {{ margin-top: 14px; border: 1px dashed var(--hairline) !important; border-radius: 10px !important; background: rgba(142, 68, 173, 0.05) !important; }} #sources-accordion .label-wrap {{ font-family: 'JetBrains Mono', monospace !important; font-size: 0.85rem !important; color: var(--purple-theme) !important; }} .source-line {{ font-family: 'JetBrains Mono', monospace; font-size: 0.82rem; color: var(--ink-teal); opacity: 0.8; padding: 3px 0; }} #safety-footer {{ text-align: center; font-size: 0.85rem; color: var(--charcoal); opacity: 0.65; padding: 18px 24px; line-height: 1.6; }} #creator-footer {{ text-align: center; padding: 22px 0 36px 0; font-family: 'Source Sans 3', sans-serif; font-size: 0.8rem; color: var(--ink-teal); opacity: 0.55; letter-spacing: 0.02em; }} footer {{ display: none !important; }} """ def format_sources_html(sources): if not sources: return "
No specific sources matched closely — answer may be general.
" lines = [] for s in sources: lines.append( f"
{s['source']}  \u00b7  page {s['page']} " f" \u00b7  relevance {s['score']:.2f}
" ) return "".join(lines) def format_answer_html(answer_text): # Convert **bold** markdown syntax to real HTML bold tags formatted = re.sub(r"\*\*(.+?)\*\*", r"\1", answer_text) formatted = formatted.replace("\n", "
") return f'
{formatted}
' def ui_ask(question, history): if not question or not question.strip(): return "", "Please enter a question above.", "" result = ask(question, k=5, max_new_tokens=384) answer_html = format_answer_html(result["answer"]) sources_html = format_sources_html(result["sources"]) return "", answer_html, sources_html with gr.Blocks(css=CUSTOM_CSS, title="TEJASWI.AI \u2014 Women's Legal Rights Assistant") as demo: gr.HTML(f"""

TEJASWI.AI

You are not alone. This is your trusted companion to navigate Indian law with absolute privacy and care. We translate complex legalities into clear, gentle guidance to help you protect your rights.

""") with gr.Column(elem_id="ask-panel"): question_box = gr.Textbox( label="Ask me anything - your rights, your voice", placeholder="e.g. What can I do if my employer ignores a harassment complaint?", lines=3, elem_id="question-input", ) with gr.Row(): ask_btn = gr.Button("Ask", elem_id="ask-button") gr.Examples( examples=EXAMPLE_QUESTIONS, inputs=question_box, label="Or try one of these", ) answer_display = gr.HTML(elem_id="answer-box", value="Your answer will appear here. Note: first response after startup may take a few minutes while the model loads.") with gr.Accordion("Sources referenced", open=False, elem_id="sources-accordion"): sources_display = gr.HTML(value="") gr.HTML(""" """) gr.HTML(f""" """) ask_btn.click( fn=ui_ask, inputs=[question_box, gr.State([])], outputs=[question_box, answer_display, sources_display], ) question_box.submit( fn=ui_ask, inputs=[question_box, gr.State([])], outputs=[question_box, answer_display, sources_display], ) if __name__ == "__main__": demo.launch()