app.py

"""
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 "<div class='source-line'>No specific sources matched closely — answer may be general.</div>"
    lines = []
    for s in sources:
        lines.append(
            f"<div class='source-line'>{s['source']} &nbsp;\u00b7&nbsp; page {s['page']} "
            f"&nbsp;\u00b7&nbsp; relevance {s['score']:.2f}</div>"
        )
    return "".join(lines)


def format_answer_html(answer_text):
    # Convert **bold** markdown syntax to real HTML bold tags
    formatted = re.sub(r"\*\*(.+?)\*\*", r"<strong>\1</strong>", answer_text)
    formatted = formatted.replace("\n", "<br>")
    return f'<div style="color: #2B2B2B; opacity: 1;">{formatted}</div>'


def ui_ask(question, history):
    if not question or not question.strip():
        return "", "<i>Please enter a question above.</i>", ""

    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"""
    <div id="hero">
        <h1>TEJASWI.AI</h1>
        <p class="tagline">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.</p>
    </div>
    <div class="book-divider"></div>
    """)

    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="<i>Your answer will appear here. Note: first response after startup may take a few minutes while the model loads.</i>")

    with gr.Accordion("Sources referenced", open=False, elem_id="sources-accordion"):
        sources_display = gr.HTML(value="")

    gr.HTML("""
    <div id="safety-footer">
        This tool provides general legal information, not legal advice, and may not
        reflect the most recent amendments. For guidance specific to your situation,
        please consult a qualified lawyer, the National Commission for Women, or your
        nearest women's helpline.
    </div>
    """)

    gr.HTML(f"""
    <div id="creator-footer">Created by {CREATOR_NAME}</div>
    """)

    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()