app.py

"""Semantic Document Search — demo for bugBug04S/legal-embed-modernbert-v2.

Doubles as a portfolio piece (preloaded contract clauses) and a sales tool
(paste any documents and search them).
"""
import spaces  # must be imported before torch / CUDA-touching packages

import hashlib

import gradio as gr
import numpy as np
from sentence_transformers import SentenceTransformer

MODEL_ID = "bugBug04S/legal-embed-modernbert-v2"

print("Loading model...")
model = SentenceTransformer(MODEL_ID)
model.max_seq_length = 256
print("Model loaded.")

SAMPLE_DOCS = """Either party may terminate this Agreement at any time, with or without cause, by providing thirty (30) days prior written notice to the other party. Upon termination, all outstanding fees shall become immediately due and payable.

The Licensee shall indemnify, defend and hold harmless the Licensor from and against any and all claims, damages, liabilities, costs and expenses arising out of or relating to the Licensee's use of the Licensed Materials.

This Agreement shall be governed by and construed in accordance with the laws of the State of Delaware, without giving effect to any choice of law or conflict of law provisions.

Each party agrees to maintain in strict confidence all Confidential Information disclosed by the other party and shall not disclose such information to any third party without prior written consent.

In no event shall either party be liable for any indirect, incidental, special, consequential or punitive damages, and the total aggregate liability of either party shall not exceed the fees paid in the twelve (12) months preceding the claim.

During the term of this Agreement and for a period of two (2) years thereafter, the Contractor shall not directly or indirectly solicit for employment any employee of the Company with whom the Contractor had material contact.

Neither party shall be liable for any failure or delay in performance resulting from causes beyond its reasonable control, including acts of God, war, terrorism, pandemic, or governmental action.

All intellectual property rights in any work product created by the Contractor in the course of performing the Services shall vest exclusively in the Company upon creation.

Any dispute arising out of or in connection with this Agreement shall be finally resolved by binding arbitration administered in accordance with the rules of the American Arbitration Association.

The Company shall pay all undisputed invoices within forty-five (45) days of receipt. Late payments shall accrue interest at the rate of 1.5% per month or the maximum rate permitted by law."""

EXAMPLES = [
    ["can we end this contract early?"],
    ["who pays if we get sued by a third party?"],
    ["what is the cap on damages?"],
    ["can I hire their staff after we finish?"],
    ["which country's law applies?"],
    ["what happens if a pandemic stops us working?"],
    ["when do invoices have to be paid?"],
]

_cache = {}


def _embed_docs(docs):
    key = hashlib.md5("||".join(docs).encode()).hexdigest()
    if key not in _cache:
        _cache[key] = model.encode(
            ["search_document: " + d for d in docs],
            normalize_embeddings=True,
            batch_size=16,
        )
        if len(_cache) > 20:          # keep memory bounded on the free tier
            _cache.pop(next(iter(_cache)))
    return _cache[key]


@spaces.GPU
def search(documents_text, query, top_k):
    docs = [d.strip() for d in (documents_text or "").split("\n\n") if d.strip()]
    if not docs:
        return "Paste some documents first — separate each passage with a blank line."
    if not (query or "").strip():
        return "Type a question to search for."

    doc_emb = _embed_docs(docs)
    q_emb = model.encode("search_query: " + query.strip(), normalize_embeddings=True)
    scores = doc_emb @ q_emb

    k = min(int(top_k), len(docs))
    order = np.argsort(-scores)[:k]

    out = [f"### Results for: *{query.strip()}*\n"]
    for rank, i in enumerate(order, 1):
        score = float(scores[i])
        bar = "█" * max(1, int(round(score * 20)))
        passage = docs[i]
        if len(passage) > 700:
            passage = passage[:700] + " …"
        out.append(
            f"**{rank}.**  match `{score:.3f}`  {bar}\n\n"
            f"> {passage}\n"
        )
    return "\n".join(out)


with gr.Blocks(title="Semantic Document Search") as demo:
    gr.Markdown(
        """
        # 🔍 Semantic Document Search

        Ask a question in plain English. This finds the passage that *answers* it —
        even when none of your words appear in the text.

        Ordinary search matches words. Ask *"can we end this contract early?"* of a
        keyword search and it looks for "end" and "early", and misses the clause that
        says *"thirty (30) days prior written notice"*. This doesn't.

        Powered by [legal-embed-modernbert-v2](https://huggingface.co/bugBug04S/legal-embed-modernbert-v2)
        — 0.844 accuracy@1 on contract clause retrieval, vs 0.644 for the base model.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            query = gr.Textbox(
                label="Your question",
                placeholder="e.g. can we end this contract early?",
                lines=2,
            )
            top_k = gr.Slider(1, 10, value=3, step=1, label="How many results")
            btn = gr.Button("Search", variant="primary", size="lg")
            gr.Examples(examples=EXAMPLES, inputs=[query], label="Try one of these")

        with gr.Column(scale=1):
            docs = gr.Textbox(
                label="Documents to search  ·  one passage per block, separated by a blank line",
                value=SAMPLE_DOCS,
                lines=14,
                max_lines=30,
            )
            gr.Markdown(
                "*Replace this with your own text — contracts, policies, manuals, "
                "FAQs, anything — and search it the same way.*"
            )

    results = gr.Markdown(label="Results")

    btn.click(search, inputs=[docs, query, top_k], outputs=results)
    query.submit(search, inputs=[docs, query, top_k], outputs=results)

if __name__ == "__main__":
    demo.launch()