Upload app.py

app.py

@@ -1,69 +1,228 @@
-import gradio as gr
-from huggingface_hub import InferenceClient
+"""
+Browser chat UI for our trained LLM.
+
+Wraps `LLM.generate()` in a Gradio ChatInterface. Streams tokens as they're
+generated. Uses ChatML formatting (same as our SFT pipeline), so it works for
+both base and SFT'd checkpoints — base models just continue more freely,
+SFT'd models actually answer.
+
+Usage (local):
+    python gradio_app.py --ckpt out/plan_b_110m/ckpt_step1000_model.pt --preset plan_b
 
+Usage (HuggingFace Spaces):
+    Place this script + model.py + configs.py + the .pt file in the Space repo.
+    Spaces auto-runs `app.py` (rename gradio_app.py → app.py).
 
-def respond(
-    message,
-    history: list[dict[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    hf_token: gr.OAuthToken,
-):
-    """
-    For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-    """
-    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
+Args:
+    --ckpt           path to checkpoint .pt file
+    --preset         which model preset (smoke, plan_b, main)
+    --tokenizer      HF repo for tokenizer (default: Llama-2 mirror)
+    --share          create a temporary 72-hour public URL (Gradio's gradio.live)
+"""
+import argparse
+import json
+import os
+from pathlib import Path
 
-    messages = [{"role": "system", "content": system_message}]
+import numpy as np
+import torch
+import gradio as gr
+from transformers import AutoTokenizer
 
-    messages.extend(history)
+from configs import PRESETS
+from model import LLM
 
-    messages.append({"role": "user", "content": message})
 
-    response = ""
+# ChatML special markers — must match sft.py
+SYS_PROMPT = "You are a helpful assistant."
+B_USR, B_ASST, EOT = "<|user|>\n", "<|assistant|>\n", "<|end|>"
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        choices = message.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
 
-        response += token
-        yield response
+class Retriever:
+    """Tiny in-memory dense retriever over the KB built by build_kb.py."""
+    def __init__(self, kb_dir):
+        kb_dir = Path(kb_dir)
+        meta = json.load(open(kb_dir / "kb.meta.json"))
+        self.chunks = [json.loads(l) for l in open(kb_dir / "kb.jsonl")]
+        self.emb = np.load(kb_dir / "kb.npy")  # already normalized
+        from sentence_transformers import SentenceTransformer
+        self.encoder = SentenceTransformer(meta["model"])
+        print(f"loaded KB: {len(self.chunks)} chunks, dim={meta['dim']}")
 
+    def topk(self, query, k=3):
+        q = self.encoder.encode([query], normalize_embeddings=True)[0]
+        scores = self.emb @ q  # cosine since both normalized
+        idx = np.argsort(-scores)[:k]
+        return [(float(scores[i]), self.chunks[i]) for i in idx]
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-chatbot = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
+
+OUT_OF_SCOPE_THRESHOLD = 0.55  # cosine similarity below this → likely off-topic for our Indian KB
+OUT_OF_SCOPE_MSG = (
+    "Sorry, this model is specialized for Indian general knowledge "
+    "(history, mythology, geography, sports, politics). I don't have "
+    "information on that topic."
 )
 
-with gr.Blocks() as demo:
-    with gr.Sidebar():
-        gr.LoginButton()
-    chatbot.render()
+
+def build_prompt(history, user_msg, retriever=None, k=5, max_ctx_chars=2400):
+    """Render conversation + user message into a ChatML prompt. If a retriever
+    is given, retrieved chunks are inlined into the user message. Returns a
+    tuple (prompt, out_of_scope) — if out_of_scope is True, the caller should
+    short-circuit and return the OUT_OF_SCOPE_MSG."""
+    parts = [f"<|system|>\n{SYS_PROMPT}{EOT}\n"]
+    for u, a in history:
+        parts.append(f"{B_USR}{u}{EOT}\n")
+        parts.append(f"{B_ASST}{a}{EOT}\n")
+
+    out_of_scope = False
+    if retriever is not None:
+        hits = retriever.topk(user_msg, k=k)
+        if hits and hits[0][0] < OUT_OF_SCOPE_THRESHOLD:
+            out_of_scope = True
+        ctx = ""
+        for score, c in hits:
+            piece = f"{c['title']}: {c['text']}\n\n"
+            remaining = max_ctx_chars - len(ctx)
+            if remaining <= 0: break
+            if len(piece) > remaining:
+                piece = piece[:remaining].rstrip() + "…\n\n"
+                ctx += piece
+                break
+            ctx += piece
+        user_block = (
+            f"Read the passage and answer the question.\n\n"
+            f"{ctx.strip()}\n\nQuestion: {user_msg}"
+        )
+    else:
+        user_block = user_msg
+
+    parts.append(f"{B_USR}{user_block}{EOT}\n{B_ASST}")
+    return "".join(parts), out_of_scope
+
+
+def make_chat_fn(model, tok, device, max_new_tokens=80, temperature=1.0,
+                 top_k=1, top_p=1.0, repetition_penalty=1.1, raw_mode=False,
+                 retriever=None):
+    # Defaults are GREEDY (top_k=1, temp=1.0, top_p=1.0) + light repetition penalty.
+    # Empirically: a 110M model with retrieved context does best with greedy
+    # decoding on factual Q&A. Sampling at any temperature destroys accuracy.
+    eos_id = tok.eos_token_id
+
+    def chat(message, history):
+        if raw_mode:
+            prompt = message
+        else:
+            prompt, oos = build_prompt(history, message, retriever=retriever)
+            if oos:
+                yield OUT_OF_SCOPE_MSG
+                return
+        ids = tok.encode(prompt, return_tensors="pt").to(device)
+        ctx = model.cfg.max_seq_len
+        if ids.shape[1] > ctx - max_new_tokens:
+            ids = ids[:, -(ctx - max_new_tokens):]
+
+        out = ids
+        is_mcq = ("(A)" in message and "(B)" in message)
+        with torch.no_grad():
+            for step in range(max_new_tokens):
+                logits, _ = model(out[:, -ctx:])
+                logits = logits[:, -1, :].float()
+
+                # Repetition penalty: shrink logits of tokens already in the output.
+                # Standard formulation from CTRL paper: divide if positive, mult if negative.
+                if repetition_penalty != 1.0:
+                    seen = out[0].tolist()
+                    for tid in set(seen[-200:]):
+                        v = logits[0, tid]
+                        logits[0, tid] = v / repetition_penalty if v > 0 else v * repetition_penalty
+
+                logits = logits / temperature
+
+                # top-k
+                if top_k:
+                    v, _ = torch.topk(logits, top_k)
+                    logits[logits < v[:, [-1]]] = -float("inf")
+
+                # top-p (nucleus): keep smallest set of tokens whose cum-prob >= top_p
+                if top_p < 1.0:
+                    sorted_logits, sorted_idx = torch.sort(logits, descending=True)
+                    cum = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
+                    mask = cum > top_p
+                    mask[..., 1:] = mask[..., :-1].clone()
+                    mask[..., 0] = False
+                    sorted_logits[mask] = -float("inf")
+                    logits = torch.full_like(logits, -float("inf")).scatter(1, sorted_idx, sorted_logits)
+
+                probs = torch.softmax(logits, dim=-1)
+                nxt = torch.multinomial(probs, 1)
+                if nxt.item() == eos_id:
+                    break
+                out = torch.cat([out, nxt], dim=1)
+
+                new_text = tok.decode(out[0, ids.shape[1]:].tolist(), skip_special_tokens=True)
+                if EOT in new_text:
+                    yield new_text.split(EOT)[0].rstrip()
+                    return
+                if is_mcq and step > 8 and (new_text.rstrip().endswith(".") or "\n" in new_text[10:]):
+                    yield new_text.split("\n")[0].rstrip()
+                    return
+                yield new_text
+
+    return chat
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", default=os.environ.get("LLM_CKPT", "model.pt"))
+    ap.add_argument("--preset", default=os.environ.get("LLM_PRESET", "plan_b"))
+    ap.add_argument("--tokenizer", default="NousResearch/Llama-2-7b-hf")
+    ap.add_argument("--share", action="store_true")
+    ap.add_argument("--raw", action="store_true")
+    ap.add_argument("--kb_dir", default=os.environ.get("LLM_KB_DIR", "kb"))
+    # On HF Spaces there are no CLI args — parse_args([]) so unknown sys.argv don't break us.
+    args, _ = ap.parse_known_args()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.float32 if device == "cpu" else torch.bfloat16
+    print(f"loading {args.ckpt} on {device}…")
+
+    mcfg, _ = PRESETS[args.preset]
+    model = LLM(mcfg).to(device, dtype=dtype)
+    state = torch.load(args.ckpt, map_location=device, weights_only=False)["model"]
+    state = {k.replace("_orig_mod.", ""): v for k, v in state.items()}
+    model.load_state_dict(state); model.eval()
+    print(f"loaded {model.num_params()/1e6:.1f}M params")
+
+    tok = AutoTokenizer.from_pretrained(args.tokenizer, use_fast=True)
+
+    retriever = None
+    if args.kb_dir and not args.raw:
+        retriever = Retriever(args.kb_dir)
+
+    chat_fn = make_chat_fn(model, tok, device, raw_mode=args.raw, retriever=retriever)
+
+    if args.raw:
+        title = "private-llm (base model) — text continuation only"
+        desc = "Base language model trained from scratch on FineWeb-Edu — give it the start of a sentence and it'll try to continue."
+        examples = [
+            "Once upon a time in a small village,",
+            "The history of India can be traced back to",
+            "Photosynthesis is the process by which",
+        ]
+    else:
+        title = "🇮🇳 private-llm — 110M Indian-GK chatbot, trained from scratch"
+        desc = "110M-param model trained from scratch + RAG over Indian Wikipedia. Indian-GK only. [GitHub](https://github.com/rsumit123/private-llm)"
+        examples = [
+            "Who is Rama's brother?",
+            "Who was the first Prime Minister of India?",
+            "What is the capital of India?",
+            "Who founded the Maratha Empire?",
+            "Tell me about Hanuman.",
+            "Who wrote the Indian National Anthem?",
+        ]
+
+    gr.ChatInterface(fn=chat_fn, title=title, description=desc, examples=examples).launch(share=args.share)
 
 
 if __name__ == "__main__":
-    demo.launch()
+    main()