app.py

import json, html, numpy as np, torch, gradio as gr
from sentence_transformers import SentenceTransformer
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
from threading import Thread


class Config:
    EMBEDDINGS_FILE = "embeddings_quality.json"
    MODEL_ID = "HuggingFaceTB/SmolLM2-135M-Instruct"
    TOP_K = 5
    SIM_THRESHOLD = 0.36
    MAX_NEW_TOKENS = 512
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


cfg = Config()


def safe_strip(x: str) -> str:
    return x.replace("\n", " ").replace("\r", " ").strip() if isinstance(x, str) else ""


def load_entries(path):
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    entries = []
    for v in data.values():
        m = v.get("metadata", {})
        title = m.get("title") or v.get("title") or ""
        definition = m.get("definition") or v.get("definition") or m.get("content") or ""
        source = m.get("source") or v.get("source") or ""
        emb = np.array(v.get("embedding", []), dtype=np.float32)
        if emb.size == 0:
            continue
        emb = emb / np.linalg.norm(emb)
        entries.append({
            "title": safe_strip(title),
            "definition": safe_strip(definition),
            "source": safe_strip(source),
            "embedding": emb
        })
    vectors = np.stack([e["embedding"] for e in entries])
    return entries, vectors


def init_models():
    # 1. Load embedding model for semantic search
    embed_model = SentenceTransformer("all-MiniLM-L6-v2", device=cfg.DEVICE)

    # 2. Load tokenizer and model for text generation
    tokenizer = AutoTokenizer.from_pretrained(cfg.MODEL_ID)

    # Add pad token if missing
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # Load model without quantization to avoid bitsandbytes issues
    model = AutoModelForCausalLM.from_pretrained(
        cfg.MODEL_ID,
        device_map="auto",
        torch_dtype=torch.float16 if cfg.DEVICE == "cuda" else torch.float32
    )

    # 3. Load entries and vectors
    entries, vectors = load_entries(cfg.EMBEDDINGS_FILE)

    return embed_model, tokenizer, model, entries, vectors


embed_model, tokenizer, model, entries, vectors = init_models()


def search_chunks(query, top_k=cfg.TOP_K, batch_size=512):
    import heapq

    # Encode query as normalized float32 vector
    qv = embed_model.encode([query], normalize_embeddings=True,
                            convert_to_numpy=True).astype("float32")[0]

    # Use a small max-heap to store the best results
    heap = []  # stores (-similarity, index)

    n = len(entries)
    for start in range(0, n, batch_size):
        end = min(start + batch_size, n)
        # Instead of dotting all vectors, dot only a slice
        sims = np.dot(vectors[start:end], qv)
        for j, s in enumerate(sims):
            if s < cfg.SIM_THRESHOLD:
                continue
            heapq.heappush(heap, (-s, start + j))
            if len(heap) > top_k:
                heapq.heappop(heap)  # maintain top_k only

    # Convert heap to sorted list (descending order)
    results = [(-s, entries[i]) for s, i in sorted(heap)]
    return [(e, float(s)) for s, e in results]


def build_context(entries, tokenizer, max_tokens=1500):
    ctx, t = [], 0
    for e in entries:
        txt = f"{e['title']}: {e['definition']}\n"
        tok = tokenizer.encode(txt, add_special_tokens=False)
        if t + len(tok) > max_tokens:
            break
        ctx.append(txt)
        t += len(tok)
    return "\n".join(ctx)


def generate_answer(question, context_entries):
    ctx_text = build_context(context_entries, tokenizer)
    prompt = f"""<|im_start|>system
You are an expert on fighting game terminology. Use only the CONTEXT below to answer the QUESTION clearly using english language and proper structure.
<|im_end|>
<|im_start|>user
CONTEXT:
{ctx_text}

QUESTION: {question}
<|im_end|>
<|im_start|>assistant
"""
    inputs = tokenizer(prompt, return_tensors="pt").to(cfg.DEVICE)
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    kwargs = dict(
        **inputs,
        max_new_tokens=cfg.MAX_NEW_TOKENS,
        temperature=0.5,
        top_p=0.9,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        streamer=streamer
    )
    Thread(target=model.generate, kwargs=kwargs).start()
    partial = ""
    for token in streamer:
        partial += token
        yield partial


def qa_pipeline(question):
    results = search_chunks(question)

    # Build top results HTML immediately
    if not results:
        top_html = "<p>No relevant entries found.</p>"
        yield top_html, "Your question is out of scope."
        return

    html_out = "<h4>Top Relevant Entries:</h4>"
    for i, (e, s) in enumerate(results, 1):
        html_out += f"<details open><summary><b>[{i}] (score: {s:.3f}) {html.escape(e['title'])}</b></summary><p>{html.escape(e['definition'][:500])}</p><p><i>{html.escape(e['source'])}</i></p></details>"

    # Yield search results immediately
    yield html_out, ""

    # Then stream the LLM response
    entries_only = [r[0] for r in results]
    partial = ""
    for token in generate_answer(question, entries_only):
        partial = token
        yield html_out, partial


with gr.Blocks(title="DustLook") as demo:
    gr.Markdown("##  DustLook\nAsk about any fighting game term.")
    q = gr.Textbox(label="Ask a question:", placeholder="e.g., What is a Roman Cancel?")
    top = gr.HTML(label="Top Matches")
    out = gr.Textbox(label="LLM Answer", lines=15, interactive=False, show_copy_button=True)
    btn = gr.Button("Search & Answer")

    # Use a single click event that streams both outputs
    btn.click(fn=qa_pipeline, inputs=q, outputs=[top, out], queue=True)

demo.queue().launch()