diff --git a/.gitignore b/.gitignore index 9e3c856fcf74f3cc3aa5c1bd5aa4f16dbd398991..83b229232bbe7dd947c05097029b2175904f9733 100644 --- a/.gitignore +++ b/.gitignore @@ -63,3 +63,6 @@ app/next-env.d.ts # Local scratch training run (checkpoint rewritten every step — never commit) runs/local_stdfix/ runs/stdfix_test/ + +# HF-mirror hygiene: snapshot_download bookkeeping, never useful to publish +data/**/.cache/ diff --git a/app/engine/generate_and_merge.py b/app/engine/generate_and_merge.py index 7750b4a977a606f0006478fc0c8ce2829eb53d70..40bfb1a2ad29d25e767e6d7f354e819dc5ec28e9 100644 --- a/app/engine/generate_and_merge.py +++ b/app/engine/generate_and_merge.py @@ -45,8 +45,34 @@ _DTYPES = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32} +def _assert_real_checkpoint(path: Path) -> None: + """Fail with an actionable message when the file is a git-LFS pointer. + + A pointer is ~130 bytes of text starting "version https://git-lfs...". + Handed to torch.load it surfaces as `Unsupported operand 118` or + `invalid load key, 'v'` -- opcode 'v', the first byte of "version" -- which + reads like a corrupt checkpoint rather than a file that was never fetched. + Any `GIT_LFS_SKIP_SMUDGE=1` clone or pull leaves checkpoints in this state. + """ + if not path.exists(): + raise FileNotFoundError(f"checkpoint not found: {path}") + if path.stat().st_size < 5000: + head = path.read_bytes()[:64] + if head.startswith(b"version https://git-lfs"): + raise RuntimeError( + f"{path} is a git-LFS pointer, not the checkpoint " + f"({path.stat().st_size} bytes). Fetch it with:\n" + f" git lfs pull --include=\"{path}\"" + ) + + def load_head(checkpoint: Path): - ckpt = torch.load(checkpoint, map_location="cpu") + _assert_real_checkpoint(Path(checkpoint)) + # weights_only=False: these checkpoints carry the run's config/args dicts, + # not just tensors, and torch>=2.6 defaults the strict unpickler on -- + # which rejects them ("Unsupported operand"). They are produced by this + # project's own training script, so loading them fully is intended. + ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False) cfg = ckpt["config"] head = MemoryLoRAHead( input_dim=cfg["input_dim"], diff --git a/app/engine/serve_fallback.py b/app/engine/serve_fallback.py index db72458b67f627dd77ad1a7a578484ba60894d91..14d478ad28bbde424464dfe725def6a11dd0e6de 100644 --- a/app/engine/serve_fallback.py +++ b/app/engine/serve_fallback.py @@ -155,6 +155,16 @@ class Engine: def stream(self, prompt: str, max_new_tokens: int, temperature: float, top_p: float, stop: list[str] | None = None, use_adapter: bool = True): + """Stream tokens, with the adapter toggled for the WHOLE operation. + + The enable/disable must wrap the entire generator -- both the worker + thread and the consumption of the streamer -- and the thread must be + joined before the context exits. PEFT's disable_adapter() flips state on + the shared model, so if the context closed while the next request was + already starting, that request would silently run with the wrong + adapter state. Putting the context inside the worker thread did exactly + that: base and adapted both came back as base. + """ from transformers import TextIteratorStreamer enc = self.tok(prompt, return_tensors="pt").to(self.device) streamer = TextIteratorStreamer( @@ -167,29 +177,35 @@ class Engine: pad_token_id=self.tok.pad_token_id or self.tok.eos_token_id, streamer=streamer, ) - def _run(): - with self._maybe_off(use_adapter): - self.model.generate(**kwargs) - threading.Thread(target=_run, daemon=True).start() - stops = self._stops(stop) - # A stop sequence can straddle two streamed pieces, so emit only the - # part of the buffer that can no longer become part of a stop string. hold = max(len(s) for s in stops) - buf, emitted = "", 0 - for piece in streamer: - buf += piece - cut, hit = self._truncate_at_stop(buf, stops) - if hit: - if len(cut) > emitted: - yield cut[emitted:] - return - safe = max(0, len(buf) - hold) - if safe > emitted: - yield buf[emitted:safe] - emitted = safe - if len(buf) > emitted: - yield buf[emitted:] + + with self._maybe_off(use_adapter): + worker = threading.Thread( + target=self.model.generate, kwargs=kwargs, daemon=True) + worker.start() + try: + # A stop sequence can straddle two streamed pieces, so emit only + # the part of the buffer that can no longer become part of one. + buf, emitted, hit = "", 0, False + for piece in streamer: + buf += piece + cut, hit = self._truncate_at_stop(buf, stops) + if hit: + if len(cut) > emitted: + yield cut[emitted:] + break + safe = max(0, len(buf) - hold) + if safe > emitted: + yield buf[emitted:safe] + emitted = safe + if not hit and len(buf) > emitted: + yield buf[emitted:] + finally: + # Drain and join so the model is idle before adapter state flips. + for _ in streamer: + pass + worker.join(timeout=120) ENGINE: Engine | None = None diff --git a/app/src/app/api/jobs/[jobId]/route.ts b/app/src/app/api/jobs/[jobId]/route.ts new file mode 100644 index 0000000000000000000000000000000000000000..143863c349f844b7d5cd18884ee31910f02c4063 --- /dev/null +++ b/app/src/app/api/jobs/[jobId]/route.ts @@ -0,0 +1,41 @@ +import { NextRequest, NextResponse } from "next/server"; +import { readStatus } from "@/lib/jobs"; +import { rmSync, existsSync } from "node:fs"; +import path from "node:path"; + +export const dynamic = "force-dynamic"; + +const WORKSPACES_DIR = path.join(process.cwd(), ".workspaces"); + +/** Delete a build: stop its inference server, then remove its workspace. */ +export async function DELETE( + _req: NextRequest, + { params }: { params: Promise<{ jobId: string }> } +) { + const { jobId } = await params; + // Reject anything that could escape the workspaces directory. + if (!/^[a-f0-9]{6,32}$/i.test(jobId)) { + return NextResponse.json({ error: "invalid job id" }, { status: 400 }); + } + const dir = path.join(WORKSPACES_DIR, jobId); + if (!existsSync(dir)) { + return NextResponse.json({ error: "not found" }, { status: 404 }); + } + + // Free the ~10GB the server holds before removing the adapter it loaded. + const status = readStatus(jobId); + const pid = status?.server?.pid; + if (pid) { + try { + process.kill(pid, "SIGTERM"); + } catch { + /* already gone */ + } + } + try { + rmSync(dir, { recursive: true, force: true }); + } catch (e) { + return NextResponse.json({ error: String(e) }, { status: 500 }); + } + return NextResponse.json({ ok: true, jobId }); +} diff --git a/app/src/app/api/jobs/[jobId]/server/route.ts b/app/src/app/api/jobs/[jobId]/server/route.ts new file mode 100644 index 0000000000000000000000000000000000000000..85338211958de1eb30ac060c41d57132977e1d70 --- /dev/null +++ b/app/src/app/api/jobs/[jobId]/server/route.ts @@ -0,0 +1,118 @@ +import { NextRequest, NextResponse } from "next/server"; +import { readStatus } from "@/lib/jobs"; +import { spawn } from "node:child_process"; +import { existsSync, readFileSync, writeFileSync } from "node:fs"; +import path from "node:path"; + +export const dynamic = "force-dynamic"; + +const APP_DIR = path.resolve(process.cwd()); +const ENGINE_DIR = path.join(APP_DIR, "engine"); +const WORKSPACES_DIR = path.join(APP_DIR, ".workspaces"); +const PYTHON = process.env.MLORA_PYTHON || "python3"; + +function statusPath(jobId: string) { + return path.join(WORKSPACES_DIR, jobId, "status.json"); +} + +/** + * Stop or restart a job's inference server. + * + * A transformers server that has been generating for a while slows down badly + * (MPS allocator churn, accumulated KV state), and the only reliable cure is a + * fresh process. Restarting reloads the same base+adapter, so the endpoint + * comes back on the same port with identical behaviour — no rebuild needed. + */ +export async function POST( + req: NextRequest, + { params }: { params: Promise<{ jobId: string }> } +) { + const { jobId } = await params; + if (!/^[a-f0-9]{6,32}$/i.test(jobId)) { + return NextResponse.json({ error: "invalid job id" }, { status: 400 }); + } + const dir = path.join(WORKSPACES_DIR, jobId); + if (!existsSync(dir)) { + return NextResponse.json({ error: "not found" }, { status: 404 }); + } + + let action = "restart"; + try { + action = (await req.json())?.action || "restart"; + } catch { + /* default */ + } + + const status = readStatus(jobId); + const oldPid = status?.server?.pid; + + // Kill first in both cases — a restart must not leave two processes fighting + // over the same port. + if (oldPid) { + try { + process.kill(oldPid, "SIGTERM"); + } catch { + /* already gone */ + } + await new Promise((r) => setTimeout(r, 1200)); + try { + process.kill(oldPid, 0); + process.kill(oldPid, "SIGKILL"); // still alive: force it + } catch { + /* exited cleanly */ + } + } + + const raw = existsSync(statusPath(jobId)) + ? JSON.parse(readFileSync(statusPath(jobId), "utf8")) + : {}; + + if (action === "stop") { + raw.state = "error"; + raw.error = "server stopped — press Restart to bring it back"; + raw.updated_at = Date.now() / 1000; + writeFileSync(statusPath(jobId), JSON.stringify(raw, null, 2)); + return NextResponse.json({ ok: true, action: "stop" }); + } + + const port = status?.server?.port || 8000; + const adapter = path.join(dir, "adapter"); + const merged = path.join(dir, "merged"); + const args = [ + path.join(ENGINE_DIR, "serve_fallback.py"), + "--port", String(port), + "--served-name", "memory-lora", + "--job", jobId, + ]; + if (existsSync(adapter)) { + args.push("--adapter", adapter); + } else if (existsSync(merged)) { + args.push("--model", merged); + } else { + return NextResponse.json( + { error: "no adapter or merged model in this workspace" }, + { status: 409 } + ); + } + if (process.env.MLORA_DEVICE) args.push("--device", process.env.MLORA_DEVICE); + + const child = spawn(PYTHON, args, { + cwd: ENGINE_DIR, + detached: true, + stdio: "ignore", + env: { ...process.env, MLORA_WORKSPACES: WORKSPACES_DIR }, + }); + child.unref(); + + // serve_fallback --job rewrites status.json with its own pid once the model + // is loaded; mark it as loading meanwhile so the UI doesn't claim it's ready. + raw.state = "running"; + raw.stage = "serve"; + raw.error = null; + raw.server = { backend: "fallback", pid: child.pid, port, mode: "lora" }; + raw.pipeline_pid = child.pid; + raw.updated_at = Date.now() / 1000; + writeFileSync(statusPath(jobId), JSON.stringify(raw, null, 2)); + + return NextResponse.json({ ok: true, action: "restart", pid: child.pid, port }); +} diff --git a/app/src/app/compare/page.tsx b/app/src/app/compare/page.tsx index a4a6ce2170647b4fbeefc55e7bf1049f79f273a3..59589b8366754217450e985111772047347c8b0a 100644 --- a/app/src/app/compare/page.tsx +++ b/app/src/app/compare/page.tsx @@ -1,270 +1,5 @@ "use client"; -import { useEffect, useRef, useState } from "react"; - -type Job = { - job_id: string; - repo_url: string; - stages: string[]; - stage: string; - state: "running" | "ready" | "error"; - error?: string | null; - server?: { backend: string; port: number; mode?: string }; -}; - -type Turn = { role: "user" | "assistant"; base?: string; adapted?: string }; - -const STAGE_LABEL: Record = { - clone: "clone", - embed: "embed", - model: "base model", - generate: "adapter", - merge: "merge", - serve: "serve", - ready: "ready", -}; - -// Questions that expose repo knowledge rather than general fluency — the base -// model can bluff a plausible answer to all of them, which is the point. -const SUGGESTED = [ - "What is the core purpose of this repository?", - "What testing framework does this repository use?", - "How is this project built and packaged?", - "What are the main modules and how do they fit together?", -]; - -export default function Compare() { - const [url, setUrl] = useState(""); - const [job, setJob] = useState(null); - const [jobs, setJobs] = useState([]); - const [turns, setTurns] = useState([]); - const [q, setQ] = useState(""); - const [busy, setBusy] = useState(false); - const scroller = useRef(null); - - useEffect(() => { - const poll = async () => { - try { - const r = await fetch("/api/build"); - const d = await r.json(); - setJobs(d.jobs || []); - setJob((cur) => - cur ? (d.jobs || []).find((j: Job) => j.job_id === cur.job_id) || cur : cur - ); - } catch { - /* ignore */ - } - }; - poll(); - const t = setInterval(poll, 2000); - return () => clearInterval(t); - }, []); - - useEffect(() => { - scroller.current?.scrollTo({ top: scroller.current.scrollHeight, behavior: "smooth" }); - }, [turns]); - - const build = async () => { - if (!url.trim()) return; - setBusy(true); - try { - const r = await fetch("/api/build", { - method: "POST", - headers: { "content-type": "application/json" }, - body: JSON.stringify({ repoUrl: url.trim() }), - }); - const d = await r.json(); - setJob({ - job_id: d.jobId, - repo_url: url.trim(), - stages: ["clone", "embed", "model", "generate", "serve", "ready"], - stage: "clone", - state: "running", - }); - setTurns([]); - setUrl(""); - } finally { - setBusy(false); - } - }; - - // Both sides hit the same server; only the model id differs, so the - // comparison is the adapter and nothing else. - const askOne = async (model: string, question: string): Promise => { - const r = await fetch("/v1/chat/completions", { - method: "POST", - headers: { "content-type": "application/json", "x-mlora-job": job!.job_id }, - body: JSON.stringify({ - model, - max_tokens: 160, - messages: [{ role: "user", content: question }], - }), - }); - if (!r.ok) return `[error ${r.status}]`; - const d = await r.json(); - return d.choices?.[0]?.message?.content?.trim() || "[empty]"; - }; - - const ask = async (question?: string) => { - const text = (question ?? q).trim(); - if (!text || !job || job.state !== "ready" || busy) return; - setBusy(true); - setQ(""); - setTurns((t) => [...t, { role: "user", base: text }, { role: "assistant" }]); - try { - const [base, adapted] = await Promise.all([ - askOne("base", text), - askOne(`memory-lora:${job.job_id}`, text), - ]); - setTurns((t) => { - const c = [...t]; - c[c.length - 1] = { role: "assistant", base, adapted }; - return c; - }); - } finally { - setBusy(false); - } - }; - - const ready = job?.state === "ready"; - const idx = job ? job.stages?.indexOf(job.stage) ?? 0 : 0; - - return ( -
-
-
-
Memory-LoRA · live comparison
-

Same model. Same prompt. One has read your repo.

-
- {job && ( - - {job.state === "ready" ? "endpoint live" : job.state} - - )} -
- -
- setUrl(e.target.value)} - onKeyDown={(e) => e.key === "Enter" && build()} - /> - -
- - {jobs.length > 0 && ( -
- Ready - {jobs - .filter((j) => j.state === "ready") - .slice(0, 4) - .map((j) => ( - - ))} -
- )} - - {job && !ready && ( -
-
- {(job.stages || []).map((s, i) => ( - - {STAGE_LABEL[s] || s} - - ))} -
- {job.state === "error" ? ( -

✗ {job.error}

- ) : ( -

- Cloning, embedding six views, and generating a LoRA adapter — about a minute. -

- )} -
- )} - - {ready && ( - <> -
-
- Frozen Gemma-4-E2B - no repo knowledge -
-
- + generated adapter - {job!.repo_url.replace("https://github.com/", "")} -
-
- -
- {turns.length === 0 && ( -
-

Ask something only this repository can answer.

-
- {SUGGESTED.map((s) => ( - - ))} -
-
- )} - {turns.map((t, i) => - t.role === "user" ? ( -
- {t.base} -
- ) : ( -
-
- {t.base ?? thinking…} -
-
- {t.adapted ?? thinking…} -
-
- ) - )} -
- -
- setQ(e.target.value)} - onKeyDown={(e) => e.key === "Enter" && ask()} - disabled={busy} - /> - -
-

- Both answers come from one resident model on the same server — the only - difference is whether the generated adapter is switched on. No repository - text is in either prompt. -

- - )} -
- ); -} +// The demo lives at "/" — this keeps "/compare" working too, so either URL +// lands on the side-by-side comparison instead of a 404. +export { default } from "../page"; diff --git a/app/src/app/globals.css b/app/src/app/globals.css index 7cf1d9ce34a25fba22bad24a782443e99833879c..eab8eaf87ff095dfb210a9a981faedb16b98de2a 100644 --- a/app/src/app/globals.css +++ b/app/src/app/globals.css @@ -162,3 +162,20 @@ label.small { color: var(--muted); font-size: 13px; } border-radius: 10px; color: var(--text); padding: 13px 15px; font-size: 15px; } .ask input:focus { outline: none; border-color: var(--accent); } .foot { color: var(--muted); font-size: 12.5px; margin: 12px 0 0; text-align: center; } +.chip { display: inline-flex; align-items: center; gap: 0; padding: 0; overflow: hidden; } +.chip-main { background: none; border: none; color: inherit; font: inherit; + padding: 5px 6px 5px 12px; cursor: pointer; font-family: var(--mono); } +.chip-main em { font-style: normal; color: var(--muted); } +.chip-x { background: none; border: none; color: var(--muted); cursor: pointer; + font-size: 15px; line-height: 1; padding: 5px 10px 6px 4px; font-weight: 400; } +.chip-x:hover { color: var(--err); } +.clear-all { background: none; border: 1px solid var(--border); color: var(--muted); + border-radius: 999px; padding: 5px 12px; font-size: 12.5px; cursor: pointer; + font-weight: 400; } +.clear-all:hover { border-color: var(--err); color: var(--err); } +button.ghost { background: transparent; border: 1px solid var(--border); + color: var(--muted); font-weight: 400; } +button.ghost:hover:not(:disabled) { border-color: var(--accent); color: var(--text); } +.head-actions { display: flex; align-items: center; gap: 8px; flex-wrap: wrap; } +button.small { padding: 5px 12px; font-size: 12.5px; border-radius: 999px; } +button.ghost.danger:hover:not(:disabled) { border-color: var(--err); color: var(--err); } diff --git a/app/src/app/page.tsx b/app/src/app/page.tsx index bb7ee6cdb9bdc773a5a27ed4d425ed7e01a14e4c..105c9166e3901c19e7db45f2a97870def6034e8c 100644 --- a/app/src/app/page.tsx +++ b/app/src/app/page.tsx @@ -1,6 +1,6 @@ "use client"; -import { useEffect, useState } from "react"; +import { useEffect, useRef, useState } from "react"; type Job = { job_id: string; @@ -9,10 +9,11 @@ type Job = { stage: string; state: "running" | "ready" | "error"; error?: string | null; - server?: { backend: string; port: number }; - endpoint?: string; + server?: { backend: string; port: number; mode?: string }; }; +type Turn = { role: "user" | "assistant"; base?: string; adapted?: string }; + const STAGE_LABEL: Record = { clone: "clone", embed: "embed", @@ -23,80 +24,37 @@ const STAGE_LABEL: Record = { ready: "ready", }; -function ConnectSnippets({ job, base }: { job: Job; base: string }) { - const [tab, setTab] = useState<"claude" | "openai" | "curl">("claude"); - const model = `memory-lora:${job.job_id}`; - const claude = `# Claude Code — point it at this repo-personalized model -export ANTHROPIC_BASE_URL="${base}" -export ANTHROPIC_API_KEY="local-demo" # any non-empty value -export ANTHROPIC_MODEL="${model}" -claude`; - const openai = `# vibe / aider / any OpenAI-compatible CLI -export OPENAI_BASE_URL="${base}/v1" -export OPENAI_API_KEY="local-demo" -export OPENAI_MODEL="${model}" -# e.g. aider --model openai/${model} -# e.g. vibe --base-url $OPENAI_BASE_URL --model ${model}`; - const curl = `curl ${base}/v1/chat/completions \\ - -H "content-type: application/json" \\ - -d '{"model":"${model}","messages":[{"role":"user","content":"What does this repo do?"}]}'`; - const text = tab === "claude" ? claude : tab === "openai" ? openai : curl; - return ( -
-
-
setTab("claude")}>Claude Code
-
setTab("openai")}>OpenAI CLIs
-
setTab("curl")}>curl
-
-
{text}
-

- Endpoint served by {job.server?.backend || "engine"} on port {job.server?.port}. The app - translates the Anthropic /v1/messages API to this model, so Claude Code works - directly. -

-
- ); -} +// Questions that expose repo knowledge rather than general fluency — the base +// model can bluff a plausible answer to all of them, which is the point. +// Short answers keep the demo snappy; the server's stop sequences cut the +// response at the end of the answer anyway, so this rarely truncates. +const MAX_TOKENS = 72; -function JobCard({ job, base }: { job: Job; base: string }) { - const idx = job.stages?.indexOf(job.stage) ?? 0; - return ( -
-
- {job.repo_url} - {job.state} -
-
- {(job.stages || []).map((s, i) => ( - - {STAGE_LABEL[s] || s} - - ))} -
- {job.state === "error" &&
✗ {job.error}
} - {job.state === "ready" && } -
- ); -} +const SUGGESTED = [ + "What is the core purpose of this repository?", + "What testing framework does this repository use?", + "How is this project built and packaged?", + "What are the main modules and how do they fit together?", +]; -export default function Home() { +export default function Compare() { const [url, setUrl] = useState(""); + const [job, setJob] = useState(null); const [jobs, setJobs] = useState([]); + const [turns, setTurns] = useState([]); + const [q, setQ] = useState(""); const [busy, setBusy] = useState(false); - const [base, setBase] = useState(""); + const scroller = useRef(null); useEffect(() => { - setBase(window.location.origin); const poll = async () => { try { const r = await fetch("/api/build"); const d = await r.json(); setJobs(d.jobs || []); + setJob((cur) => + cur ? (d.jobs || []).find((j: Job) => j.job_id === cur.job_id) || cur : cur + ); } catch { /* ignore */ } @@ -106,56 +64,322 @@ export default function Home() { return () => clearInterval(t); }, []); - const submit = async () => { + useEffect(() => { + scroller.current?.scrollTo({ top: scroller.current.scrollHeight, behavior: "smooth" }); + }, [turns]); + + const build = async () => { if (!url.trim()) return; setBusy(true); try { - await fetch("/api/build", { + const r = await fetch("/api/build", { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ repoUrl: url.trim() }), }); + const d = await r.json(); + setJob({ + job_id: d.jobId, + repo_url: url.trim(), + stages: ["clone", "embed", "model", "generate", "serve", "ready"], + stage: "clone", + state: "running", + }); + setTurns([]); setUrl(""); } finally { setBusy(false); } }; + // Both sides hit the same server; only the model id differs, so the + // comparison is the adapter and nothing else. + // + // Streamed, and deliberately sequential: one model serves both sides, so the + // server holds a lock and concurrent requests would queue anyway. Streaming + // is what makes it feel immediate -- tokens appear as they are produced + // instead of after the full answer. + const askOne = async ( + model: string, + question: string, + onDelta: (chunk: string) => void + ): Promise => { + const r = await fetch("/v1/chat/completions", { + method: "POST", + headers: { "content-type": "application/json", "x-mlora-job": job!.job_id }, + body: JSON.stringify({ + model, + max_tokens: MAX_TOKENS, + stream: true, + messages: [{ role: "user", content: question }], + }), + }); + if (!r.ok || !r.body) { + onDelta(`[error ${r.status}]`); + return; + } + const reader = r.body.getReader(); + const dec = new TextDecoder(); + let buf = ""; + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + buf += dec.decode(value, { stream: true }); + const lines = buf.split("\n"); + buf = lines.pop() || ""; + for (const line of lines) { + const t = line.trim(); + if (!t.startsWith("data:")) continue; + const payload = t.slice(5).trim(); + if (payload === "[DONE]") continue; + try { + const d = JSON.parse(payload); + const piece = d.choices?.[0]?.delta?.content; + if (piece) onDelta(piece); + } catch { + /* partial frame */ + } + } + } + }; + + const ask = async (question?: string) => { + const text = (question ?? q).trim(); + if (!text || !job || job.state !== "ready" || busy) return; + setBusy(true); + setQ(""); + setTurns((t) => [ + ...t, + { role: "user", base: text }, + { role: "assistant", base: "", adapted: "" }, + ]); + const push = (side: "base" | "adapted") => (chunk: string) => + setTurns((t) => { + const c = [...t]; + const last = { ...c[c.length - 1] }; + last[side] = (last[side] || "") + chunk; + c[c.length - 1] = last; + return c; + }); + try { + await askOne("base", text, push("base")); + await askOne(`memory-lora:${job.job_id}`, text, push("adapted")); + } finally { + setBusy(false); + } + }; + + const clearChat = () => setTurns([]); + + // A long-lived transformers server degrades badly (allocator churn on MPS); + // a fresh process reloads the same base+adapter on the same port. + const serverAction = async (action: "stop" | "restart") => { + if (!job) return; + setBusy(true); + try { + await fetch(`/api/jobs/${job.job_id}/server`, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ action }), + }); + const r = await fetch("/api/build"); + setJobs((await r.json()).jobs || []); + } finally { + setBusy(false); + } + }; + + const deleteBuild = async (jobId: string) => { + await fetch(`/api/jobs/${jobId}`, { method: "DELETE" }); + if (job?.job_id === jobId) { + setJob(null); + setTurns([]); + } + const r = await fetch("/api/build"); + setJobs((await r.json()).jobs || []); + }; + + const clearAll = async () => { + await Promise.all(jobs.map((j) => fetch(`/api/jobs/${j.job_id}`, { method: "DELETE" }))); + setJob(null); + setTurns([]); + setJobs([]); + }; + + const ready = job?.state === "ready"; + const idx = job ? job.stages?.indexOf(job.stage) ?? 0 : 0; + return ( -
-

Memory-LoRA

-

- Paste a git repository. A hypernetwork reads a 6-view embedding of the codebase and{" "} - emits a LoRA adapter for Gemma in one forward pass — no fine-tuning. The adapter is - merged into the base model and served behind an OpenAI + Anthropic compatible endpoint, so - your coding CLI talks to a model that already knows this repo. -

- -
-
- setUrl(e.target.value)} - onKeyDown={(e) => e.key === "Enter" && submit()} - /> - +
+
+
+
Memory-LoRA · live comparison
+

Same model. Same prompt. One has read your repo.

-

- First build downloads the base model (~10 GB) and takes a few minutes; later builds - reuse it. -

+ {job && ( +
+ + {job.state === "ready" ? "endpoint live" : job.state} + + + +
+ )} +
+ +
+ setUrl(e.target.value)} + onKeyDown={(e) => e.key === "Enter" && build()} + /> +
- {jobs.length === 0 && ( -

No builds yet. Paste a repo URL above to start one.

+ {jobs.length > 0 && ( +
+ Ready + {jobs.slice(0, 5).map((j) => ( + + + + + ))} + {jobs.length > 0 && ( + + )} +
+ )} + + {job && !ready && ( +
+
+ {(job.stages || []).map((s, i) => ( + + {STAGE_LABEL[s] || s} + + ))} +
+ {job.state === "error" ? ( +

✗ {job.error}

+ ) : ( +

+ Cloning, embedding six views, and generating a LoRA adapter — about a minute. +

+ )} +
+ )} + + {ready && ( + <> +
+
+ Frozen Gemma-4-E2B + no repo knowledge +
+
+ + generated adapter + {job!.repo_url.replace("https://github.com/", "")} +
+
+ +
+ {turns.length === 0 && ( +
+

Ask something only this repository can answer.

+
+ {SUGGESTED.map((s) => ( + + ))} +
+
+ )} + {turns.map((t, i) => + t.role === "user" ? ( +
+ {t.base} +
+ ) : ( +
+
+ {t.base ? t.base : } +
+
+ {t.adapted ? t.adapted : } +
+
+ ) + )} +
+ +
+ setQ(e.target.value)} + onKeyDown={(e) => e.key === "Enter" && ask()} + disabled={busy} + /> + + {turns.length > 0 && ( + + )} +
+

+ Both answers come from one resident model on the same server — the only + difference is whether the generated adapter is switched on. No repository + text is in either prompt. +

+ )} - {jobs.map((j) => ( - - ))}
); } diff --git a/app/src/app/serve/page.tsx b/app/src/app/serve/page.tsx new file mode 100644 index 0000000000000000000000000000000000000000..bb7ee6cdb9bdc773a5a27ed4d425ed7e01a14e4c --- /dev/null +++ b/app/src/app/serve/page.tsx @@ -0,0 +1,161 @@ +"use client"; + +import { useEffect, useState } from "react"; + +type Job = { + job_id: string; + repo_url: string; + stages: string[]; + stage: string; + state: "running" | "ready" | "error"; + error?: string | null; + server?: { backend: string; port: number }; + endpoint?: string; +}; + +const STAGE_LABEL: Record = { + clone: "clone", + embed: "embed", + model: "base model", + generate: "adapter", + merge: "merge", + serve: "serve", + ready: "ready", +}; + +function ConnectSnippets({ job, base }: { job: Job; base: string }) { + const [tab, setTab] = useState<"claude" | "openai" | "curl">("claude"); + const model = `memory-lora:${job.job_id}`; + const claude = `# Claude Code — point it at this repo-personalized model +export ANTHROPIC_BASE_URL="${base}" +export ANTHROPIC_API_KEY="local-demo" # any non-empty value +export ANTHROPIC_MODEL="${model}" +claude`; + const openai = `# vibe / aider / any OpenAI-compatible CLI +export OPENAI_BASE_URL="${base}/v1" +export OPENAI_API_KEY="local-demo" +export OPENAI_MODEL="${model}" +# e.g. aider --model openai/${model} +# e.g. vibe --base-url $OPENAI_BASE_URL --model ${model}`; + const curl = `curl ${base}/v1/chat/completions \\ + -H "content-type: application/json" \\ + -d '{"model":"${model}","messages":[{"role":"user","content":"What does this repo do?"}]}'`; + const text = tab === "claude" ? claude : tab === "openai" ? openai : curl; + return ( +
+
+
setTab("claude")}>Claude Code
+
setTab("openai")}>OpenAI CLIs
+
setTab("curl")}>curl
+
+
{text}
+

+ Endpoint served by {job.server?.backend || "engine"} on port {job.server?.port}. The app + translates the Anthropic /v1/messages API to this model, so Claude Code works + directly. +

+
+ ); +} + +function JobCard({ job, base }: { job: Job; base: string }) { + const idx = job.stages?.indexOf(job.stage) ?? 0; + return ( +
+
+ {job.repo_url} + {job.state} +
+
+ {(job.stages || []).map((s, i) => ( + + {STAGE_LABEL[s] || s} + + ))} +
+ {job.state === "error" &&
✗ {job.error}
} + {job.state === "ready" && } +
+ ); +} + +export default function Home() { + const [url, setUrl] = useState(""); + const [jobs, setJobs] = useState([]); + const [busy, setBusy] = useState(false); + const [base, setBase] = useState(""); + + useEffect(() => { + setBase(window.location.origin); + const poll = async () => { + try { + const r = await fetch("/api/build"); + const d = await r.json(); + setJobs(d.jobs || []); + } catch { + /* ignore */ + } + }; + poll(); + const t = setInterval(poll, 2000); + return () => clearInterval(t); + }, []); + + const submit = async () => { + if (!url.trim()) return; + setBusy(true); + try { + await fetch("/api/build", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ repoUrl: url.trim() }), + }); + setUrl(""); + } finally { + setBusy(false); + } + }; + + return ( +
+

Memory-LoRA

+

+ Paste a git repository. A hypernetwork reads a 6-view embedding of the codebase and{" "} + emits a LoRA adapter for Gemma in one forward pass — no fine-tuning. The adapter is + merged into the base model and served behind an OpenAI + Anthropic compatible endpoint, so + your coding CLI talks to a model that already knows this repo. +

+ +
+
+ setUrl(e.target.value)} + onKeyDown={(e) => e.key === "Enter" && submit()} + /> + +
+

+ First build downloads the base model (~10 GB) and takes a few minutes; later builds + reuse it. +

+
+ + {jobs.length === 0 && ( +

No builds yet. Paste a repo URL above to start one.

+ )} + {jobs.map((j) => ( + + ))} +
+ ); +} diff --git a/app/src/app/v1/[...path]/route.ts b/app/src/app/v1/[...path]/route.ts index 2a8e9960f374c854cede0ffae18c02a9574c61a6..1f43b394bfff230dc9d0c1d14f07a667e0a98e10 100644 --- a/app/src/app/v1/[...path]/route.ts +++ b/app/src/app/v1/[...path]/route.ts @@ -11,6 +11,11 @@ export const dynamic = "force-dynamic"; const SERVED_MODEL = "memory-lora"; const API_KEY = process.env.MLORA_API_KEY || ""; +/** True when the caller explicitly asked for the un-adapted base model. */ +function isBaseModel(model: unknown): boolean { + return typeof model === "string" && model.trim().toLowerCase() === "base"; +} + function authOk(req: NextRequest): boolean { if (!API_KEY) return true; const auth = req.headers.get("authorization"); @@ -53,12 +58,15 @@ async function proxyPassthrough( ); } - // Force the served model name so arbitrary client model strings work. + // Normalize the model field so arbitrary client model strings work — but + // NEVER collapse "base" into the served name. The upstream server selects + // frozen-vs-adapted from this field, so rewriting it unconditionally made the + // side-by-side comparison silently return the adapted model on BOTH sides. let body = rawBody; if (rawBody) { try { const j = JSON.parse(rawBody); - j.model = SERVED_MODEL; + if (!isBaseModel(j.model)) j.model = SERVED_MODEL; body = JSON.stringify(j); } catch { /* leave as-is */ @@ -96,7 +104,8 @@ async function handleMessages(req: NextRequest, rawBody: string): Promise=2.6 defaults the strict unpickler on -- + # which rejects them ("Unsupported operand"). They are produced by this + # project's own training script, so loading them fully is intended. + ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False) tokenizer = AutoTokenizer.from_pretrained(model_name) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token diff --git a/scripts/train_memory_lora.py b/scripts/train_memory_lora.py index 9dee693187da25f5b621da2ca3a03b11fb6a22d9..f80e74e4e7e6ac0894a534283c1f25be8c0912f1 100644 --- a/scripts/train_memory_lora.py +++ b/scripts/train_memory_lora.py @@ -493,7 +493,11 @@ def main() -> None: print(f" fitted input standardization over {len(train_docs)} train docs", flush=True) if args.resume_from: - ckpt = torch.load(args.resume_from, map_location=device) + # weights_only=False: these checkpoints carry the run's config/args dicts, + # not just tensors, and torch>=2.6 defaults the strict unpickler on -- + # which rejects them ("Unsupported operand"). They are produced by this + # project's own training script, so loading them fully is intended. + ckpt = torch.load(args.resume_from, map_location=device, weights_only=False) # strict=False: checkpoints written before input standardization # existed carry no input_mean/input_std, so the stats fitted above are # kept. A checkpoint that does carry them overwrites the fresh fit,