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 = "No relevant entries found.
"
yield top_html, "Your question is out of scope."
return
html_out = "Top Relevant Entries:
"
for i, (e, s) in enumerate(results, 1):
html_out += f"[{i}] (score: {s:.3f}) {html.escape(e['title'])}
{html.escape(e['definition'][:500])}
{html.escape(e['source'])}