Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| import io | |
| import json | |
| import math | |
| import tempfile | |
| import time | |
| import torch | |
| import gradio as gr | |
| import spaces | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| from analysis import token_surprisal, surprisal_html | |
| from styles import STYLES, HELD_OUT, NEUTRAL | |
| MODEL_REPO = "OpenPeerAI/My-Name-is-Riemann-Prerelease" | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO) | |
| tokenizer.pad_token = tokenizer.eos_token | |
| model = AutoModelForCausalLM.from_pretrained( | |
| MODEL_REPO, trust_remote_code=True, dtype=torch.bfloat16, low_cpu_mem_usage=True | |
| ) | |
| model.eval() | |
| N_PARAMS = sum(p.numel() for p in model.parameters()) | |
| def _gpu(): | |
| if model.device.type != "cuda": | |
| model.to("cuda") | |
| def mean_nll(texts): | |
| total, n = 0.0, 0 | |
| for t in texts: | |
| r = token_surprisal(model, tokenizer, t, max_length=256) | |
| total += r["nll"] * r["n_tokens"] | |
| n += r["n_tokens"] | |
| return total / max(1, n) | |
| def gate_plot(tuner): | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt | |
| raw = tuner.controller.raw.detach().float().cpu().numpy() | |
| layers = tuner.controller.layer_indices.detach().cpu().tolist() | |
| counts = {} | |
| for l in layers: | |
| counts[l] = counts.get(l, 0) + 1 | |
| fig, axes = plt.subplots(1, 2, figsize=(10, 3.4)) | |
| axes[0].hist(raw, bins=40, color="#4c6ef5") | |
| axes[0].set_title("Learned log-gate values s (layer output x exp(s))") | |
| axes[1].bar(list(counts.keys()), list(counts.values()), color="#20c997") | |
| axes[1].set_title("Gates per decoder layer") | |
| axes[1].set_xlabel("layer") | |
| fig.tight_layout() | |
| buf = io.BytesIO() | |
| fig.savefig(buf, format="png", dpi=120) | |
| plt.close(fig) | |
| buf.seek(0) | |
| return buf | |
| def fit( | |
| target_style, gates, steps, lr, max_log_gate, probe_text, progress=gr.Progress() | |
| ): | |
| from ntkmirror import ForwardFineTuner, Example | |
| from PIL import Image | |
| _gpu() | |
| train = [Example(prompt="", completion=t) for t in STYLES[target_style]] | |
| eval_sets = {**HELD_OUT, "Neutral prose": NEUTRAL} | |
| progress(0.05, desc="Scoring base model on every style") | |
| before = {k: mean_nll(v) for k, v in eval_sets.items()} | |
| probe_before = token_surprisal(model, tokenizer, probe_text, max_length=256) | |
| tuner = ForwardFineTuner( | |
| model, tokenizer, gates=int(gates), max_log_gate=float(max_log_gate) | |
| ) | |
| progress(0.2, desc=f"Fitting gates on {len(train)} {target_style} sentences") | |
| t0 = time.perf_counter() | |
| tuner.fit( | |
| train, | |
| steps=int(steps), | |
| lr=float(lr), | |
| batch_size=8, | |
| max_length=128, | |
| verbose=False, | |
| ) | |
| elapsed = time.perf_counter() - t0 | |
| progress(0.8, desc="Scoring with controller") | |
| tuner.controller.attach() | |
| after = {k: mean_nll(v) for k, v in eval_sets.items()} | |
| probe_after = token_surprisal(model, tokenizer, probe_text, max_length=256) | |
| tuner.save( | |
| tmp_path := tempfile.NamedTemporaryFile(suffix=".pt", delete=False).name, | |
| metadata={"demo": "ntk-trainer-space", "style": target_style}, | |
| ) | |
| tuner.controller.remove() | |
| n_gates = int(tuner.controller.raw.numel()) | |
| rows = [] | |
| for k in eval_sets: | |
| d = after[k] - before[k] | |
| mark = " (target)" if k == target_style else "" | |
| rows.append( | |
| f"| {k}{mark} | {math.exp(before[k]):.1f} | {math.exp(after[k]):.1f} | {d:+.3f} | {'lower' if d < -0.005 else ('higher' if d > 0.005 else 'about the same')} |" | |
| ) | |
| table = ( | |
| "| Held-out style | Perplexity, base | Perplexity, with controller | Change in mean surprisal (nats) | Direction |\n|---|---|---|---|---|\n" | |
| + "\n".join(rows) | |
| ) | |
| target_gain = before[target_style] - after[target_style] | |
| others = [before[k] - after[k] for k in eval_sets if k != target_style] | |
| selectivity = target_gain - (sum(others) / len(others)) | |
| probe_html = ( | |
| surprisal_html(probe_before, "Base weights:") | |
| + "<hr/>" | |
| + surprisal_html(probe_after, f"With controller fitted on {target_style}:") | |
| + f"<p>Mean surprisal on this text: {probe_before['nll']:.3f} to {probe_after['nll']:.3f} nats ({probe_after['nll']-probe_before['nll']:+.3f}).</p>" | |
| ) | |
| summary = { | |
| "base_model": MODEL_REPO, | |
| "base_parameters": N_PARAMS, | |
| "trainable_gates": n_gates, | |
| "trainable_fraction": n_gates / N_PARAMS, | |
| "target_style": target_style, | |
| "train_sentences": len(train), | |
| "steps": int(steps), | |
| "fit_seconds": round(elapsed, 1), | |
| "heldout_nll_before": {k: round(v, 4) for k, v in before.items()}, | |
| "heldout_nll_after": {k: round(v, 4) for k, v in after.items()}, | |
| "target_gain_nats": round(target_gain, 4), | |
| "mean_other_gain_nats": round(sum(others) / len(others), 4), | |
| "selectivity_nats": round(selectivity, 4), | |
| } | |
| report = ( | |
| f"**Base model:** {MODEL_REPO}, {N_PARAMS/1e9:.1f}B parameters, frozen throughout.\n\n" | |
| f"**Controller:** {n_gates:,} gates ({100*n_gates/N_PARAMS:.5f}% of the model), fitted on {len(train)} {target_style} sentences in {int(steps)} steps ({elapsed:.1f}s). The saved controller file is a few kilobytes.\n\n" | |
| f"**Result on held-out sentences the gates never saw:**\n\n{table}\n\n" | |
| f"The controller lowered surprisal on held-out **{target_style}** by **{target_gain:+.3f} nats** while moving the other styles by {sum(others)/len(others):+.3f} on average, a selectivity of **{selectivity:+.3f} nats**. " | |
| "A positive selectivity means the gates learned something specific to the target style rather than a general shift. " | |
| "This is the same mechanism that tuned the 13.1B Prerelease toward its question-answering data: the weights never change, the controller re-weights what the model already represents." | |
| ) | |
| img = Image.open(gate_plot(tuner)) | |
| return report, probe_html, img, json.dumps(summary, indent=2), tmp_path | |
| DESCRIPTION = f""" | |
| # NTK Trainer Demo | |
| Live demonstration of the | |
| [OpenPeer-AI-NTK-Trainer](https://huggingface.co/OpenPeerAI/OpenPeer-AI-NTK-Trainer): | |
| adapting a frozen language model by fitting a small set of **signed log-gates** on its layer outputs, | |
| with the fitting objective derived from neural tangent kernel theory, instead of training LoRA adapters or | |
| full weights. | |
| The base model here is the **{N_PARAMS/1e9:.1f}B My Name is Riemann Prerelease**, the model this trainer | |
| was actually used on. Pick a target writing style; the space fits a controller on forty sentences in that | |
| style, then scores held-out sentences from **every** style, before and after. A working controller lowers | |
| the model's surprise on the target style and leaves the others alone. You also see, token by token, how the | |
| gates changed the model's view of any text you paste, and you can download the controller. | |
| Gates re-weight what the model already knows. They do not add facts or teach new output formats, so this | |
| demo measures what they change rather than asking them to produce a specific answer. | |
| """ | |
| with gr.Blocks(title="NTK Trainer Demo") as demo: | |
| gr.Markdown(DESCRIPTION) | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| style = gr.Dropdown( | |
| choices=list(STYLES), | |
| value="Legal and formal", | |
| label="Target style to fit the controller on", | |
| ) | |
| with gr.Accordion( | |
| "Training sentences for this style (read-only preview)", open=False | |
| ): | |
| preview = gr.Textbox( | |
| value="\n".join(STYLES["Legal and formal"][:8]) + "\n...", | |
| lines=9, | |
| interactive=False, | |
| show_label=False, | |
| ) | |
| style.change(lambda s: "\n".join(STYLES[s][:8]) + "\n...", style, preview) | |
| probe = gr.Textbox( | |
| label="Probe text for the token map (any sentence)", | |
| value="The parties agree that any notice shall be delivered in writing to the registered address.", | |
| lines=2, | |
| ) | |
| with gr.Column(scale=1): | |
| gates = gr.Slider(256, 8192, value=2048, step=256, label="Gates to fit") | |
| steps = gr.Slider(20, 300, value=120, step=10, label="Fitting steps") | |
| lr = gr.Slider(1e-3, 3e-2, value=1e-2, step=1e-3, label="Learning rate") | |
| max_log_gate = gr.Slider( | |
| 0.02, 0.5, value=0.2, step=0.02, label="Max |log gate|" | |
| ) | |
| run = gr.Button("Fit controller and measure", variant="primary") | |
| report = gr.Markdown() | |
| probe_out = gr.HTML(label="Token-level effect on the probe text") | |
| with gr.Row(): | |
| plot = gr.Image(label="Learned gates", type="pil") | |
| with gr.Column(): | |
| summary = gr.Code(label="Summary JSON", language="json") | |
| download = gr.File(label="Download controller (.pt)") | |
| run.click( | |
| fit, | |
| [style, gates, steps, lr, max_log_gate, probe], | |
| [report, probe_out, plot, summary, download], | |
| ) | |
| gr.Markdown( | |
| "OpenPeer-AI-NTK-Trainer is released under MIT. Depends on [ntkmirror](https://github.com/leochlon/ntkmirror). " | |
| "Maintained by OpenPeer AI." | |
| ) | |
| demo.queue(max_size=8).launch() | |