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_all.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 load_entries(path):
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)

    entries, embeddings = [], []

    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 = v.get("embedding", None)

        if not emb or not isinstance(emb, list): 
            continue

        emb = np.asarray(emb, dtype=np.float32)
        norm = np.linalg.norm(emb)
        if norm == 0 or np.isnan(norm): 
            continue
        emb /= norm

        entries.append({
            "title": safe_strip(title),
            "definition": safe_strip(definition[:1000]),  # clip huge text
            "source": safe_strip(source),
        })
        embeddings.append(emb)


    vectors = np.memmap("vectors.dat", dtype=np.float32, mode="w+", shape=(len(embeddings), len(embeddings[0])))
    for i, e in enumerate(embeddings):
        vectors[i] = e
    vectors.flush()

    print(f"Loaded {len(entries)} entries with shape {vectors.shape}")
    return entries, np.memmap("vectors.dat", dtype=np.float32, mode="r", shape=vectors.shape)


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

def search_chunks(query, top_k=cfg.TOP_K, batch_size=512):
    """Memory-efficient cosine similarity search."""
    # 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):
    """Unified pipeline that streams search results first, then LLM answer."""
    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="Fighting Game Glossary QA") as demo:
    gr.Markdown("## 🎮 Fighting Game Glossary QA\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()