Publish openpeerllm-demo

README.md

@@ -1,6 +1,6 @@
 ---
 title: OpenPeerLLM Demo
-emoji: 🕸️
+emoji: 🕸️
 colorFrom: blue
 colorTo: green
 sdk: gradio

analysis.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+import html
+import math
+import torch
+import torch.nn.functional as F
+
+
+@torch.no_grad()
+def token_surprisal(model, tokenizer, text: str, max_length: int = 1024):
+    ids = tokenizer(
+        text, return_tensors="pt", truncation=True, max_length=max_length
+    ).input_ids.to(model.device)
+    logits = model(input_ids=ids).logits[0, :-1].float()
+    targets = ids[0, 1:]
+    logprobs = F.log_softmax(logits, dim=-1)
+    tok_lp = logprobs.gather(1, targets[:, None]).squeeze(1)
+    ranks = (logprobs > tok_lp[:, None]).sum(dim=1) + 1
+    tokens = [tokenizer.decode([t]) for t in ids[0].tolist()]
+    nll = -tok_lp.mean().item()
+    return {
+        "tokens": tokens,
+        "surprisal": [-x for x in tok_lp.tolist()],
+        "rank": ranks.tolist(),
+        "nll": nll,
+        "ppl": math.exp(nll),
+        "n_tokens": int(targets.numel()),
+    }
+
+
+def surprisal_html(result: dict, title: str = "") -> str:
+    toks, surp = result["tokens"], result["surprisal"]
+    cells = [f"<span style='color:#888'>{html.escape(toks[0])}</span>"]
+    for t, s in zip(toks[1:], surp):
+        p = min(1.0, s / 12.0)
+        r, g = int(255 * p), int(200 * (1 - p))
+        cells.append(
+            f"<span title='surprisal {s:.2f} nats' style='background:rgba({r},{g},80,0.35);padding:1px 2px;border-radius:3px'>{html.escape(t)}</span>"
+        )
+    head = (
+        f"<div style='margin-bottom:6px'><b>{html.escape(title)}</b> perplexity <b>{result['ppl']:.1f}</b> over {result['n_tokens']} tokens (mean surprisal {result['nll']:.2f} nats)</div>"
+        if title
+        else ""
+    )
+    return (
+        head
+        + "<div style='font-family:ui-monospace,monospace;line-height:1.9;white-space:pre-wrap;word-break:break-word'>"
+        + "".join(cells)
+        + "</div><div style='color:#777;font-size:0.85em;margin-top:4px'>green = expected, red = surprising; hover a token for its surprisal</div>"
+    )
+
+
+@torch.no_grad()
+def greedy_continue(
+    model,
+    tokenizer,
+    prompt: str,
+    max_new_tokens: int,
+    repetition_penalty: float = 1.2,
+    temperature: float = 0.0,
+    top_p: float = 0.9,
+    max_context: int = 1024,
+):
+    ids = tokenizer(
+        prompt,
+        return_tensors="pt",
+        truncation=True,
+        max_length=max_context - max_new_tokens,
+    ).input_ids.to(model.device)
+    out = model.generate(
+        ids,
+        max_new_tokens=max_new_tokens,
+        do_sample=temperature > 0,
+        temperature=max(temperature, 1e-3),
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        pad_token_id=tokenizer.eos_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+    )
+    return tokenizer.decode(out[0][ids.shape[1] :], skip_special_tokens=True)
+
+
+def contrast_rows(model, tokenizer, samples: list[tuple[str, str]]):
+    rows = []
+    for label, text in samples:
+        r = token_surprisal(model, tokenizer, text)
+        rows.append([label, round(r["ppl"], 1), round(r["nll"], 3), text])
+    return rows
+
+
+def md_table(rows, header=("Text", "Perplexity", "Mean surprisal (nats)")):
+    out = ["| " + " | ".join(header) + " |", "|" + "---|" * len(header)]
+    for r in rows:
+        out.append(f"| {r[0]} | {r[1]} | {r[2]} |")
+    return "\n".join(out)
+
+
+WIKI_FEWSHOT = (
+    "Paris is the capital and most populous city of France. It is located on the Seine river in the north of the country.\n\n"
+    "Tokyo is the capital of Japan and one of the most populous cities in the world. It lies on the eastern coast of Honshu.\n\n"
+)
+
+QA_FEWSHOT = (
+    "Q: What is the capital of France?\nA: Paris.\n\n"
+    "Q: Which river flows through Cairo?\nA: The Nile.\n\n"
+    "Q: What is the largest planet in the Solar System?\nA: Jupiter.\n\n"
+)
+
+CONTRAST_SET = [
+    (
+        "Fluent English (encyclopedic)",
+        "The Amazon River in South America is the largest river by discharge volume of water in the world.",
+    ),
+    (
+        "Fluent English (narrative)",
+        "She closed the book, turned off the lamp, and listened to the rain against the window.",
+    ),
+    (
+        "Grammatical but false",
+        "The Amazon River is located in northern Canada and flows into the Arctic Ocean near Greenland.",
+    ),
+    (
+        "Word salad (same words shuffled)",
+        "world the in water of volume discharge by river largest the is America South in River Amazon The",
+    ),
+    ("Random words", "purple elephant calculus whisper nine umbrella fjord lantern"),
+    ("Python", "def add(a, b):\n    return a + b"),
+]

app.py

@@ -1,18 +1,28 @@
 from __future__ import annotations
-import threading
 import torch
 import gradio as gr
 import spaces
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
-from transformers import AutoTokenizer, TextIteratorStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from configuration_riemann import RiemannConfig
 from modeling_riemann import RiemannForCausalLM
+from analysis import (
+    token_surprisal,
+    surprisal_html,
+    greedy_continue,
+    contrast_rows,
+    md_table,
+    WIKI_FEWSHOT,
+    CONTRAST_SET,
+)
 
 OPENPEER_REPO = "OpenPeerAI/OpenPeerLLM"
+RIEMANN_REPO = "OpenPeerAI/My-Name-is-Riemann-Prerelease"
 
 tokenizer = AutoTokenizer.from_pretrained("gpt2")
 tokenizer.pad_token = tokenizer.eos_token
+
 config = RiemannConfig(
     vocab_size=50257,
     hidden_size=768,
@@ -21,113 +31,165 @@ config = RiemannConfig(
     intermediate_size=3072,
     max_position_embeddings=1024,
 )
-model = RiemannForCausalLM(config)
-weights_path = hf_hub_download(OPENPEER_REPO, "model.safetensors")
-state = load_file(weights_path)
-missing, unexpected = model.load_state_dict(state, strict=False)
-LOAD_REPORT = f"Loaded {len(state)} tensors from {OPENPEER_REPO}/model.safetensors. Missing: {len(missing)}. Unexpected: {len(unexpected)}."
-model.eval()
-N_PARAMS = sum(p.numel() for p in model.parameters())
+openpeer = RiemannForCausalLM(config)
+_state = load_file(hf_hub_download(OPENPEER_REPO, "model.safetensors"))
+_missing, _unexpected = openpeer.load_state_dict(_state, strict=False)
+openpeer.eval()
+LOAD_REPORT = f"{len(_state)} tensors loaded from {OPENPEER_REPO}, {len(_missing)} missing, {len(_unexpected)} unexpected."
+N_OPENPEER = sum(p.numel() for p in openpeer.parameters())
+
+_riemann = {"model": None}
+
+
+def riemann():
+    if _riemann["model"] is None:
+        m = AutoModelForCausalLM.from_pretrained(
+            RIEMANN_REPO,
+            trust_remote_code=True,
+            dtype=torch.bfloat16,
+            low_cpu_mem_usage=True,
+        )
+        m.eval()
+        _riemann["model"] = m
+    return _riemann["model"]
+
+
+def _gpu(m):
+    if m.device.type != "cuda":
+        m.to("cuda")
 
 
 @spaces.GPU(duration=60)
-def generate(prompt, max_new_tokens, temperature, top_p, repetition_penalty):
-    model.to("cuda")
-    inputs = tokenizer(
-        prompt,
-        return_tensors="pt",
-        truncation=True,
-        max_length=1024 - int(max_new_tokens),
-    ).to("cuda")
-    streamer = TextIteratorStreamer(
-        tokenizer, skip_prompt=True, skip_special_tokens=True
+def complete(subject, opening, max_new, compare):
+    prompt_tail = f"{subject.strip()} {opening.strip()}".strip()
+    _gpu(openpeer)
+    a = greedy_continue(
+        openpeer,
+        tokenizer,
+        WIKI_FEWSHOT + prompt_tail,
+        int(max_new),
+        repetition_penalty=1.3,
     )
-    kwargs = dict(
-        **inputs,
-        streamer=streamer,
-        max_new_tokens=int(max_new_tokens),
-        do_sample=temperature > 0,
-        temperature=max(float(temperature), 1e-3),
-        top_p=float(top_p),
-        repetition_penalty=float(repetition_penalty),
-        pad_token_id=tokenizer.eos_token_id,
+    out = (
+        f"### OpenPeerLLM ({N_OPENPEER/1e6:.0f}M, the starting rung)\n{prompt_tail}{a}"
     )
-    thread = threading.Thread(target=model.generate, kwargs=kwargs)
-    thread.start()
-    out = ""
-    for piece in streamer:
-        out += piece
-        yield prompt + out
-    thread.join()
-
-
-@spaces.GPU(duration=30)
-def score(text):
-    model.to("cuda")
-    ids = tokenizer(text, return_tensors="pt", truncation=True, max_length=1024).to(
-        "cuda"
+    if compare:
+        r = riemann()
+        _gpu(r)
+        b = greedy_continue(
+            r,
+            tokenizer,
+            WIKI_FEWSHOT + prompt_tail,
+            int(max_new),
+            repetition_penalty=1.25,
+            max_context=2048,
+        )
+        out += f"\n\n### My Name is Riemann Prerelease (13.1B, six rungs later)\n{prompt_tail}{b}"
+    return out
+
+
+@spaces.GPU(duration=60)
+def analyse(text, compare):
+    _gpu(openpeer)
+    a = token_surprisal(openpeer, tokenizer, text)
+    out = surprisal_html(a, f"OpenPeerLLM ({N_OPENPEER/1e6:.0f}M):")
+    if compare:
+        r = riemann()
+        _gpu(r)
+        b = token_surprisal(r, tokenizer, text, max_length=2048)
+        out += "<hr/>" + surprisal_html(b, "My Name is Riemann Prerelease (13.1B):")
+        out += f"<p><b>Perplexity {a['ppl']:.1f} to {b['ppl']:.1f}</b> across the growth ladder on this text.</p>"
+    return out
+
+
+@spaces.GPU(duration=90)
+def contrast(compare):
+    _gpu(openpeer)
+    rows_a = contrast_rows(openpeer, tokenizer, CONTRAST_SET)
+    if not compare:
+        return md_table([[r[0], r[1], r[2]] for r in rows_a])
+    r = riemann()
+    _gpu(r)
+    rows_b = contrast_rows(r, tokenizer, CONTRAST_SET)
+    header = ("Text", "OpenPeerLLM 163M", "Riemann Prerelease 13.1B")
+    rows = [[ra[0], ra[1], rb[1]] for ra, rb in zip(rows_a, rows_b)]
+    return (
+        md_table(rows, header)
+        + "\n\nPerplexity on each string. Lower is better on fluent text; the gap between fluent text and noise "
+        "should widen as the model improves. OpenPeerLLM is the 163M checkpoint that every later rung inherits from."
     )
-    with torch.no_grad():
-        out = model(input_ids=ids["input_ids"], labels=ids["input_ids"])
-    n = ids["input_ids"].shape[1]
-    return f"Tokens: {n}\nMean NLL: {out.loss.item():.4f}\nPerplexity: {torch.exp(out.loss).item():.2f}"
 
 
 DESCRIPTION = f"""
 # OpenPeerLLM Demo
 
-[OpenPeerAI/OpenPeerLLM](https://huggingface.co/{OPENPEER_REPO}) is a decentralised-training
-language model by Andrew Magdy Kamal Nassief (DOI [10.57967/hf/6469](https://doi.org/10.57967/hf/6469)).
-This space loads its published `model.safetensors` ({N_PARAMS/1e6:.1f}M parameters, 12 layers,
-hidden size 768, GPT-2 vocabulary) into a faithful reimplementation of its architecture
-and runs it on ZeroGPU.
-
-{LOAD_REPORT}
+[OpenPeerAI/OpenPeerLLM](https://huggingface.co/{OPENPEER_REPO}) is a decentralised-training language model
+by Andrew Magdy Kamal Nassief (DOI [10.57967/hf/6469](https://doi.org/10.57967/hf/6469)). This space loads its
+published `model.safetensors` ({N_OPENPEER/1e6:.0f}M parameters) into a faithful reimplementation of its
+architecture, with the causal attention mask the reference forward pass omits. {LOAD_REPORT}
 
-OpenPeerLLM is the starting rung of the
-[My Name is Riemann](https://huggingface.co/OpenPeerAI/My-Name-is-Riemann-Prerelease) growth ladder.
-The reimplementation here adds a causal attention mask, which the reference forward pass omits and
-which is required for coherent autoregressive generation.
+OpenPeerLLM is the **starting rung** of the
+[My Name is Riemann](https://huggingface.co/{RIEMANN_REPO}) growth ladder. Every tab lets you put the
+13.1B Prerelease next to it on the same input, so the effect of six rungs of growth and training is visible
+directly. Both are base models: they complete and score text, they do not chat.
 """
 
 with gr.Blocks(title="OpenPeerLLM Demo") as demo:
     gr.Markdown(DESCRIPTION)
-    with gr.Tab("Generate"):
-        prompt = gr.Textbox(
-            label="Prompt",
-            lines=4,
-            value="The theory of decentralised machine learning says that",
+    with gr.Tab("Complete an article"):
+        gr.Markdown(
+            "Two short example articles are prepended so the model knows the register."
         )
         with gr.Row():
-            max_new = gr.Slider(8, 256, value=96, step=8, label="Max new tokens")
-            temp = gr.Slider(
-                0.0, 1.5, value=0.8, step=0.05, label="Temperature (0 = greedy)"
+            subject = gr.Textbox(label="Subject", value="The Nile")
+            opening = gr.Textbox(
+                label="Opening clause",
+                value="is a major north-flowing river in northeastern Africa. It",
+            )
+        with gr.Row():
+            max_new = gr.Slider(16, 120, value=64, step=8, label="Max new tokens")
+            cmp1 = gr.Checkbox(
+                value=True, label="Also run My Name is Riemann Prerelease (13.1B)"
             )
-            top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
-            rep = gr.Slider(1.0, 1.5, value=1.2, step=0.05, label="Repetition penalty")
-        out = gr.Textbox(label="Output", lines=10)
-        gr.Button("Generate", variant="primary").click(
-            generate, [prompt, max_new, temp, top_p, rep], out
+        out1 = gr.Markdown()
+        gr.Button("Complete", variant="primary").click(
+            complete, [subject, opening, max_new, cmp1], out1
         )
         gr.Examples(
             [
-                ["Peer-to-peer networks allow"],
-                ["def fibonacci(n):"],
-                ["In 1859, Bernhard Riemann"],
+                ["Mount Kilimanjaro", "is a dormant volcano in Tanzania. It"],
+                [
+                    "The Rosetta Stone",
+                    "is a stele inscribed with three versions of a decree issued in",
+                ],
             ],
-            inputs=prompt,
+            inputs=[subject, opening],
+        )
+    with gr.Tab("Analyse text"):
+        gr.Markdown(
+            "Each token is coloured by how surprised the model was to see it; green is expected, red is surprising."
+        )
+        a_text = gr.Textbox(
+            label="Text",
+            lines=4,
+            value="The Great Pyramid of Giza is the oldest of the Seven Wonders of the Ancient World.",
+        )
+        cmp2 = gr.Checkbox(
+            value=True, label="Also analyse with My Name is Riemann Prerelease (13.1B)"
+        )
+        out2 = gr.HTML()
+        gr.Button("Analyse", variant="primary").click(analyse, [a_text, cmp2], out2)
+    with gr.Tab("Fluency test"):
+        gr.Markdown(
+            "Six fixed strings from fluent English down to random words, scored by perplexity."
         )
-    with gr.Tab("Score text"):
-        txt = gr.Textbox(
-            label="Text to score",
-            lines=6,
-            value="Decentralized training distributes model updates across peers.",
+        cmp3 = gr.Checkbox(
+            value=True, label="Compare with My Name is Riemann Prerelease (13.1B)"
         )
-        res = gr.Textbox(label="Negative log-likelihood and perplexity", lines=3)
-        gr.Button("Score", variant="primary").click(score, txt, res)
+        out3 = gr.Markdown()
+        gr.Button("Run fluency test", variant="primary").click(contrast, cmp3, out3)
     gr.Markdown(
-        "OpenPeerLLM is released under OPNL, MIT and CC BY 4.0. "
-        "This demo is maintained by OpenPeer AI."
+        "OpenPeerLLM is released under OPNL, MIT and CC BY 4.0. Maintained by OpenPeer AI."
     )
 
 demo.queue(max_size=16).launch()