ExtraltodeusandClaude Opus 4.8
d9773394e3
The fit corpus was limited to three hardcoded choices (wikitext, Semantic-Harmless, mixed), and any other id was rejected. Now any HuggingFace dataset id works, and any number of them can be ticked to fit on an equal-parts mix, shuffled. n_prompts now counts training SEQUENCES (what the fit iterates over) instead of source rows: each dataset is packed up to its quota, so the number entered is exactly what runs, regardless of the dataset. The fixed dropdown becomes a checkable dataset library persisted in localStorage. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1371 lines
45 KiB
Python
1371 lines
45 KiB
Python
import asyncio
|
||
import json
|
||
import logging
|
||
import mimetypes
|
||
import os
|
||
import re
|
||
import threading
|
||
from pathlib import Path
|
||
|
||
from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
|
||
|
||
mimetypes.add_type("text/javascript", ".js")
|
||
mimetypes.add_type("text/css", ".css")
|
||
from fastapi.responses import Response
|
||
from fastapi.staticfiles import StaticFiles
|
||
from pydantic import BaseModel
|
||
|
||
import config
|
||
from core import editing, registry
|
||
from core.ablation import Interventions
|
||
from core.neighbors import TokenNeighbors
|
||
from core import fitting
|
||
from core.fitting import FitManager
|
||
from core.gpus import gpu_stats
|
||
from core.lens_manager import LensManager
|
||
from core.model_manager import (
|
||
ModelManager,
|
||
_resolve_revision,
|
||
resolve_local_dir,
|
||
resolve_source,
|
||
)
|
||
from core.store import Store
|
||
|
||
manager = ModelManager()
|
||
lens_manager = LensManager()
|
||
store = Store()
|
||
fit_manager = FitManager()
|
||
interventions = Interventions()
|
||
neighbors = TokenNeighbors()
|
||
app = FastAPI(title="J-Wash")
|
||
|
||
_ws_locks = {}
|
||
_loop_holder = {}
|
||
|
||
|
||
def _valid_devices():
|
||
"""Accepted devices = "auto" + one cuda:N per GPU actually present.
|
||
Adaptive: no longer assumes the personal 2-GPU (cuda:0/cuda:1) setup."""
|
||
try:
|
||
n = len(gpu_stats())
|
||
except Exception:
|
||
n = 0
|
||
return {"auto"} | {f"cuda:{i}" for i in range(n)}
|
||
|
||
|
||
class _QuietPolling(logging.Filter):
|
||
"""Drops the access lines from the UI polling (GET /api/status every 2 s)."""
|
||
|
||
def filter(self, record):
|
||
return "GET /api/status " not in record.getMessage()
|
||
|
||
|
||
@app.on_event("startup")
|
||
async def _on_startup():
|
||
_loop_holder["loop"] = asyncio.get_running_loop()
|
||
logging.getLogger("uvicorn.access").addFilter(_QuietPolling())
|
||
|
||
|
||
async def _ws_send(ws, text):
|
||
lock = _ws_locks.get(ws)
|
||
if lock is None:
|
||
return
|
||
async with lock:
|
||
await ws.send_text(text)
|
||
|
||
|
||
def _broadcast_fit(state):
|
||
loop = _loop_holder.get("loop")
|
||
if loop is None:
|
||
return
|
||
payload = json.dumps({"type": "fit_progress", "fit": state})
|
||
for ws in list(_ws_locks):
|
||
asyncio.run_coroutine_threadsafe(_ws_send(ws, payload), loop)
|
||
|
||
|
||
fit_manager.on_progress = _broadcast_fit
|
||
|
||
# concurrent HF downloads: one state per repo_id
|
||
_downloads = {}
|
||
_downloads_lock = threading.Lock()
|
||
|
||
|
||
class LoadRequest(BaseModel):
|
||
model_id: str
|
||
dtype: str = config.DEFAULT_DTYPE
|
||
quant: str | None = None
|
||
device: str = config.DEFAULT_DEVICE
|
||
|
||
|
||
class DownloadRequest(BaseModel):
|
||
repo_id: str
|
||
|
||
|
||
class DeleteModelRequest(BaseModel):
|
||
model_id: str
|
||
|
||
|
||
class LensLoadRequest(BaseModel):
|
||
repo_id: str | None = None
|
||
filename: str = "lens.pt"
|
||
revision: str | None = None
|
||
path: str | None = None
|
||
layers: list[int] | None = None
|
||
k: int = 8
|
||
|
||
|
||
class LensLayersRequest(BaseModel):
|
||
layers: list[int]
|
||
k: int | None = None
|
||
|
||
|
||
class PinRequest(BaseModel):
|
||
gen_id: int
|
||
token_ids: list[int]
|
||
|
||
|
||
class ConversationPatch(BaseModel):
|
||
title: str | None = None
|
||
tags: list[str] | None = None
|
||
|
||
|
||
class InterventionRequest(BaseModel):
|
||
token_id: int
|
||
mode: str = "scale"
|
||
factor: float = 0.0
|
||
replacement_id: int | None = None
|
||
layers: list[int] | None = None
|
||
|
||
|
||
_NAME_RE = re.compile(r"[\w][\w.\- ]*", re.UNICODE)
|
||
|
||
|
||
def _safe_name(name):
|
||
"""Validate a user-supplied file/folder name (presets, exports): plain
|
||
names only — no separators, no traversal."""
|
||
name = (name or "").strip()
|
||
if not name or ".." in name or not _NAME_RE.fullmatch(name):
|
||
raise HTTPException(
|
||
422, f"invalid name {name!r}: letters, digits, spaces, . - _ only"
|
||
)
|
||
return name
|
||
|
||
|
||
class InterventionPatch(BaseModel):
|
||
factor: float | None = None
|
||
layers: list[int] | None = None
|
||
enabled: bool | None = None
|
||
token_id: int | None = None
|
||
replacement_id: int | None = None
|
||
mode: str | None = None # scale | replace
|
||
|
||
|
||
class InterventionsScale(BaseModel):
|
||
scale: float | None = None
|
||
mode: str | None = None
|
||
|
||
|
||
class ExportRequest(BaseModel):
|
||
format: str = "layers"
|
||
name: str
|
||
|
||
|
||
class FitRequest(BaseModel):
|
||
model_id: str
|
||
dtype: str = "bf16"
|
||
quant: str | None = None
|
||
n_prompts: int = 100
|
||
datasets: list[str] = [fitting.DATASET_WIKITEXT] # any HF ids; several = equal-parts mix
|
||
devices: list[str] = ["cuda:0"]
|
||
name: str | None = None
|
||
dim_batch: int | None = None
|
||
max_seq_len: int = 128
|
||
source_layers: list[int] | None = None
|
||
continue_from: str | None = None
|
||
|
||
|
||
@app.get("/api/models")
|
||
def api_models():
|
||
return {"models": manager.list_models()}
|
||
|
||
|
||
@app.get("/api/status")
|
||
def api_status():
|
||
return {
|
||
"loaded": manager.meta,
|
||
"busy": manager.busy,
|
||
"lens": lens_manager.meta,
|
||
"gpus": gpu_stats(),
|
||
"downloads": list(_downloads.values()),
|
||
"convert": _convert_state,
|
||
"fit": fit_manager.state,
|
||
"gguf": dict(_gguf_state),
|
||
"interventions": interventions.summary(),
|
||
"interventions_scale": interventions.global_scale,
|
||
"interventions_mode": interventions.mode,
|
||
}
|
||
|
||
|
||
@app.post("/api/load")
|
||
async def api_load(req: LoadRequest):
|
||
if req.dtype not in config.DTYPES:
|
||
raise HTTPException(422, f"invalid dtype: {req.dtype}")
|
||
if req.quant is not None and req.quant not in config.QUANTS:
|
||
raise HTTPException(422, f"invalid quant: {req.quant}")
|
||
if req.device not in _valid_devices():
|
||
raise HTTPException(422, f"invalid device: {req.device}")
|
||
if manager.busy:
|
||
raise HTTPException(409, f"busy: {manager.busy}")
|
||
try:
|
||
return await asyncio.to_thread(
|
||
manager.load, req.model_id, req.dtype, req.quant, req.device
|
||
)
|
||
except Exception as exc:
|
||
raise HTTPException(500, str(exc))
|
||
|
||
|
||
@app.post("/api/models/delete")
|
||
async def api_delete_model(req: DeleteModelRequest):
|
||
if manager.busy:
|
||
raise HTTPException(409, f"busy: {manager.busy}")
|
||
if manager.meta and manager.meta.get("model_id") == req.model_id:
|
||
raise HTTPException(409, "unload this model before deleting it")
|
||
from core.model_manager import delete_model
|
||
try:
|
||
return await asyncio.to_thread(delete_model, req.model_id)
|
||
except ValueError as exc:
|
||
raise HTTPException(400, str(exc))
|
||
except Exception as exc:
|
||
raise HTTPException(500, str(exc))
|
||
|
||
|
||
# --- user settings (Options tab) --------------------------------------------
|
||
SETTINGS_PATH = config.DATA_DIR / "settings.json"
|
||
SETTINGS_DEFAULTS = {
|
||
"default_quant": "", # '', 'int8' or 'nf4' — preselected in the Model tab
|
||
"auto_layer_radius": 2, # editor: peak ± radius when auto-selecting layers
|
||
"chat_markdown": True, # render assistant replies as markdown
|
||
"hf_cache": "", # HF cache dir — applied at startup (--hf-cache wins)
|
||
"llamacpp_dir": "", # llama.cpp folder → enables the direct GGUF export
|
||
}
|
||
|
||
|
||
def read_settings():
|
||
try:
|
||
data = json.loads(SETTINGS_PATH.read_text(encoding="utf-8"))
|
||
except Exception:
|
||
data = {}
|
||
return {**SETTINGS_DEFAULTS, **{k: v for k, v in data.items() if k in SETTINGS_DEFAULTS}}
|
||
|
||
|
||
class SettingsPatch(BaseModel):
|
||
default_quant: str | None = None
|
||
auto_layer_radius: int | None = None
|
||
chat_markdown: bool | None = None
|
||
hf_cache: str | None = None
|
||
llamacpp_dir: str | None = None
|
||
|
||
|
||
@app.get("/api/settings")
|
||
def api_settings():
|
||
return read_settings()
|
||
|
||
|
||
@app.patch("/api/settings")
|
||
def api_settings_patch(req: SettingsPatch):
|
||
if req.default_quant is not None and req.default_quant not in ("", "int8", "nf4"):
|
||
raise HTTPException(422, f"invalid quant: {req.default_quant}")
|
||
current = read_settings()
|
||
for key, value in req.model_dump(exclude_none=True).items():
|
||
if key == "auto_layer_radius":
|
||
value = max(0, min(8, int(value)))
|
||
current[key] = value
|
||
SETTINGS_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||
SETTINGS_PATH.write_text(
|
||
json.dumps(current, ensure_ascii=False, indent=1), encoding="utf-8"
|
||
)
|
||
return current
|
||
|
||
|
||
class RegisterModelRequest(BaseModel):
|
||
path: str
|
||
|
||
|
||
@app.post("/api/models/register")
|
||
def api_models_register(req: RegisterModelRequest):
|
||
"""Add a model folder to the available list (no copy — just remembered)."""
|
||
from core.model_manager import register_model_dir
|
||
try:
|
||
return register_model_dir(req.path)
|
||
except ValueError as exc:
|
||
raise HTTPException(400, str(exc))
|
||
|
||
|
||
@app.post("/api/models/unregister")
|
||
def api_models_unregister(req: RegisterModelRequest):
|
||
"""Forget a registered entry; the model files are left untouched."""
|
||
from core.model_manager import unregister_model_dir
|
||
try:
|
||
return unregister_model_dir(req.path)
|
||
except ValueError as exc:
|
||
raise HTTPException(404, str(exc))
|
||
|
||
|
||
@app.post("/api/unload")
|
||
async def api_unload():
|
||
if manager.busy:
|
||
raise HTTPException(409, f"busy: {manager.busy}")
|
||
lens_manager.unload()
|
||
neighbors.reset()
|
||
return await asyncio.to_thread(manager.unload)
|
||
|
||
|
||
@app.post("/api/lens/load")
|
||
async def api_lens_load(req: LensLoadRequest):
|
||
if not req.repo_id and not req.path:
|
||
raise HTTPException(422, "repo_id or path required")
|
||
if manager.busy:
|
||
raise HTTPException(409, f"busy: {manager.busy}")
|
||
try:
|
||
return await asyncio.to_thread(
|
||
lens_manager.load,
|
||
manager,
|
||
repo_id=req.repo_id,
|
||
filename=req.filename,
|
||
revision=req.revision,
|
||
path=req.path,
|
||
layers=req.layers,
|
||
k=req.k,
|
||
)
|
||
except Exception as exc:
|
||
raise HTTPException(500, str(exc))
|
||
|
||
|
||
@app.post("/api/lens/unload")
|
||
def api_lens_unload():
|
||
return lens_manager.unload()
|
||
|
||
|
||
@app.post("/api/lens/layers")
|
||
def api_lens_layers(req: LensLayersRequest):
|
||
if manager.busy:
|
||
raise HTTPException(409, f"busy: {manager.busy}")
|
||
try:
|
||
return lens_manager.set_layers(manager, req.layers, k=req.k)
|
||
except ValueError as exc:
|
||
raise HTTPException(422, str(exc))
|
||
|
||
|
||
@app.get("/api/interventions")
|
||
def api_interventions():
|
||
return {"rules": interventions.summary()}
|
||
|
||
|
||
@app.post("/api/interventions")
|
||
def api_interventions_add(req: InterventionRequest):
|
||
if manager.hf_model is None or lens_manager.lens is None:
|
||
raise HTTPException(422, "model and lens required")
|
||
try:
|
||
return {
|
||
"rules": interventions.add(
|
||
lens_manager,
|
||
manager.jl,
|
||
token_id=req.token_id,
|
||
mode=req.mode,
|
||
factor=req.factor,
|
||
replacement_id=req.replacement_id,
|
||
layers=req.layers,
|
||
)
|
||
}
|
||
except ValueError as exc:
|
||
raise HTTPException(422, str(exc))
|
||
|
||
|
||
@app.patch("/api/interventions/{rule_id}")
|
||
def api_interventions_patch(rule_id: int, req: InterventionPatch):
|
||
needs_dirs = any(
|
||
x is not None for x in (req.layers, req.token_id, req.replacement_id, req.mode)
|
||
)
|
||
try:
|
||
return {
|
||
"rules": interventions.update(
|
||
rule_id,
|
||
factor=req.factor,
|
||
layers=req.layers,
|
||
enabled=req.enabled,
|
||
token_id=req.token_id,
|
||
replacement_id=req.replacement_id,
|
||
mode=req.mode,
|
||
lens_manager=lens_manager if needs_dirs else None,
|
||
jl=manager.jl if needs_dirs else None,
|
||
)
|
||
}
|
||
except ValueError as exc:
|
||
raise HTTPException(404, str(exc))
|
||
|
||
|
||
@app.patch("/api/interventions")
|
||
def api_interventions_scale(req: InterventionsScale):
|
||
if req.scale is not None:
|
||
interventions.set_scale(req.scale)
|
||
try:
|
||
if req.mode is not None:
|
||
if (
|
||
req.mode in ("readthrough", "exact")
|
||
and manager.meta is not None
|
||
and manager.meta.get("rebase_supported") is False
|
||
):
|
||
raise HTTPException(
|
||
422,
|
||
"read projection unavailable on this architecture (write "
|
||
"norms, Gemma style) — use \"abliteration\" for pure weights",
|
||
)
|
||
interventions.set_mode(req.mode)
|
||
except ValueError as exc:
|
||
raise HTTPException(422, str(exc))
|
||
return {
|
||
"scale": interventions.global_scale,
|
||
"mode": interventions.mode,
|
||
}
|
||
|
||
|
||
@app.delete("/api/interventions/{rule_id}")
|
||
def api_interventions_remove(rule_id: int):
|
||
return {"rules": interventions.remove(rule_id)}
|
||
|
||
|
||
@app.delete("/api/interventions")
|
||
def api_interventions_clear():
|
||
return {"rules": interventions.remove()}
|
||
|
||
|
||
@app.get("/api/presets")
|
||
def api_presets():
|
||
return {"presets": editing.list_presets()}
|
||
|
||
|
||
@app.post("/api/presets/{name}")
|
||
def api_presets_save(name: str):
|
||
name = _safe_name(name)
|
||
rules = interventions.summary()
|
||
if not rules:
|
||
raise HTTPException(422, "no active intervention to save")
|
||
return editing.save_preset(
|
||
name, rules, manager.meta.get("model_id") if manager.meta else None,
|
||
scale=interventions.global_scale,
|
||
)
|
||
|
||
|
||
@app.post("/api/presets/{name}/apply")
|
||
def api_presets_apply(name: str):
|
||
name = _safe_name(name)
|
||
if manager.hf_model is None or lens_manager.lens is None:
|
||
raise HTTPException(422, "model and lens required")
|
||
try:
|
||
preset = editing.load_preset(name)
|
||
except ValueError as exc:
|
||
raise HTTPException(404, str(exc))
|
||
warnings = []
|
||
if preset.get("model_id") and manager.meta and preset["model_id"] != manager.meta["model_id"]:
|
||
warnings.append(
|
||
f"preset saved for {preset['model_id']}, loaded model: {manager.meta['model_id']}"
|
||
)
|
||
rules = None
|
||
for rule in preset.get("rules", []):
|
||
try:
|
||
rules = interventions.add(
|
||
lens_manager,
|
||
manager.jl,
|
||
token_id=rule["token_id"],
|
||
mode=rule["mode"],
|
||
factor=rule["factor"],
|
||
replacement_id=rule.get("replacement_id"),
|
||
layers=rule.get("layers"),
|
||
enabled=rule.get("enabled", True),
|
||
)
|
||
except ValueError as exc:
|
||
warnings.append(f"rule {rule.get('token')!r} skipped: {exc}")
|
||
if preset.get("scale") is not None:
|
||
interventions.set_scale(preset["scale"])
|
||
return {
|
||
"rules": rules or interventions.summary(),
|
||
"scale": interventions.global_scale,
|
||
"warnings": warnings,
|
||
}
|
||
|
||
|
||
@app.delete("/api/presets/{name}")
|
||
def api_presets_delete(name: str):
|
||
editing.delete_preset(_safe_name(name))
|
||
return {"ok": True}
|
||
|
||
|
||
@app.post("/api/edit/export")
|
||
async def api_edit_export(req: ExportRequest):
|
||
req.name = _safe_name(req.name)
|
||
if manager.hf_model is None:
|
||
raise HTTPException(422, "no model loaded")
|
||
rules = interventions.active_rules_full()
|
||
if not rules:
|
||
raise HTTPException(422, "no active intervention to export (rules disabled or without layers?)")
|
||
if req.format not in ("layers", "lora", "full"):
|
||
raise HTTPException(422, f"unknown format: {req.format}")
|
||
source_dir = resolve_local_dir(manager.meta["model_id"])
|
||
mode = interventions.mode
|
||
kwargs = {}
|
||
if mode in ("readthrough", "exact"):
|
||
export_fn = editing.export_rebase
|
||
kwargs["exact"] = mode == "exact"
|
||
elif mode == "abliteration":
|
||
export_fn = editing.export_abliteration
|
||
else:
|
||
raise HTTPException(
|
||
422,
|
||
"export requires a pure-weights mode: switch to \"read projection\" "
|
||
"(or \"global projection\" on write-norm architectures) — per-layer "
|
||
"steering does not bake faithfully",
|
||
)
|
||
try:
|
||
return await asyncio.to_thread(
|
||
export_fn,
|
||
rules,
|
||
manager.jl,
|
||
manager.meta,
|
||
fmt=req.format,
|
||
name=req.name,
|
||
source_dir=source_dir,
|
||
scale=interventions.global_scale,
|
||
**kwargs,
|
||
)
|
||
except ValueError as exc:
|
||
raise HTTPException(422, str(exc))
|
||
except Exception as exc:
|
||
raise HTTPException(500, str(exc))
|
||
finally:
|
||
# a full export builds ~2× the model in RAM (source + edited tensors);
|
||
# on failure, force a collection so those copies don't linger
|
||
import gc
|
||
gc.collect()
|
||
|
||
|
||
# --- direct GGUF export (via a user-provided llama.cpp folder) ---------------
|
||
# Two stages: (1) bake the full HF checkpoint into data/edits/<name>/hf — kept
|
||
# as a CACHE so several GGUF types can be exported without re-baking — then
|
||
# (2) convert_hf_to_gguf.py (+ llama-quantize for quantized types) in a
|
||
# background thread, progress polled through /api/status.
|
||
_gguf_state = {"state": "idle", "name": None, "step": None, "error": None, "result": None}
|
||
|
||
GGUF_BASE_TYPES = ("bf16", "f16")
|
||
GGUF_QUANT_TYPES = ("q8_0", "q6_k", "q5_k_m", "q4_k_m", "q3_k_m")
|
||
|
||
|
||
class GGUFExportRequest(BaseModel):
|
||
name: str
|
||
gguf_type: str = "q4_k_m"
|
||
|
||
|
||
def _llamacpp_paths():
|
||
"""(convert_py, quantize_exe, gguf_py) from the configured llama.cpp dir."""
|
||
root = read_settings().get("llamacpp_dir") or ""
|
||
root = Path(root).expanduser() if root else None
|
||
if not root or not root.is_dir():
|
||
raise ValueError(
|
||
"llama.cpp folder not set — configure it in the Options tab to "
|
||
"enable the direct GGUF export"
|
||
)
|
||
convert = root / "convert_hf_to_gguf.py"
|
||
if not convert.exists():
|
||
raise ValueError(f"convert_hf_to_gguf.py not found in {root}")
|
||
quantize = None
|
||
for cand in ("llama-quantize", "llama-quantize.exe"):
|
||
for sub in (".", "bin", "build/bin"):
|
||
p = root / sub / cand
|
||
if p.exists():
|
||
quantize = p
|
||
break
|
||
if quantize:
|
||
break
|
||
gguf_py = root / "gguf-py"
|
||
return convert, quantize, (gguf_py if gguf_py.is_dir() else None)
|
||
|
||
|
||
def _gguf_worker(name, gguf_type, hf_dir, convert, quantize, gguf_py):
|
||
import subprocess
|
||
import sys
|
||
try:
|
||
out_dir = editing.EDITS_DIR / name
|
||
base_type = gguf_type if gguf_type in GGUF_BASE_TYPES else "bf16"
|
||
base_gguf = out_dir / f"{name}-{base_type}.gguf"
|
||
env = dict(os.environ)
|
||
if gguf_py is not None: # vendored gguf package inside the llama.cpp repo
|
||
env["PYTHONPATH"] = str(gguf_py) + os.pathsep + env.get("PYTHONPATH", "")
|
||
if not base_gguf.exists():
|
||
_gguf_state.update(step=f"converting to {base_type}")
|
||
proc = subprocess.run(
|
||
[sys.executable, "-X", "utf8", str(convert), str(hf_dir),
|
||
"--outfile", str(base_gguf), "--outtype", base_type],
|
||
capture_output=True, text=True, env=env,
|
||
)
|
||
if proc.returncode != 0:
|
||
raise RuntimeError(f"convert_hf_to_gguf failed: {proc.stderr[-2000:]}")
|
||
result_path = base_gguf
|
||
if gguf_type not in GGUF_BASE_TYPES:
|
||
if quantize is None:
|
||
raise RuntimeError(
|
||
"llama-quantize not found in the llama.cpp folder — only "
|
||
"bf16/f16 exports are possible"
|
||
)
|
||
_gguf_state.update(step=f"quantizing to {gguf_type}")
|
||
result_path = out_dir / f"{name}-{gguf_type}.gguf"
|
||
proc = subprocess.run(
|
||
[str(quantize), str(base_gguf), str(result_path), gguf_type],
|
||
capture_output=True, text=True,
|
||
)
|
||
if proc.returncode != 0:
|
||
result_path.unlink(missing_ok=True)
|
||
raise RuntimeError(f"llama-quantize failed: {proc.stderr[-2000:]}")
|
||
_gguf_state.update(
|
||
state="done", step=None, error=None,
|
||
result={
|
||
"gguf": str(result_path),
|
||
"size_bytes": result_path.stat().st_size,
|
||
"hf_cache": str(hf_dir),
|
||
},
|
||
)
|
||
except Exception as exc:
|
||
_gguf_state.update(state="error", step=None, error=str(exc))
|
||
|
||
|
||
@app.post("/api/edit/export-gguf")
|
||
async def api_edit_export_gguf(req: GGUFExportRequest):
|
||
req.name = _safe_name(req.name)
|
||
if _gguf_state["state"] == "running":
|
||
raise HTTPException(409, "a GGUF export is already in progress")
|
||
if req.gguf_type not in GGUF_BASE_TYPES + GGUF_QUANT_TYPES:
|
||
raise HTTPException(422, f"unknown GGUF type: {req.gguf_type}")
|
||
try:
|
||
convert, quantize, gguf_py = _llamacpp_paths()
|
||
except ValueError as exc:
|
||
raise HTTPException(422, str(exc))
|
||
if req.gguf_type not in GGUF_BASE_TYPES and quantize is None:
|
||
raise HTTPException(422, "llama-quantize not found — pick bf16 or f16")
|
||
|
||
hf_dir = editing.EDITS_DIR / req.name / "hf"
|
||
baked = "reused"
|
||
if not (hf_dir / "config.json").exists():
|
||
# no cached checkpoint: bake one from the ACTIVE rules (same path as a
|
||
# plain full export)
|
||
if manager.hf_model is None:
|
||
raise HTTPException(422, "no model loaded (and no cached checkpoint for this name)")
|
||
rules = interventions.active_rules_full()
|
||
if not rules:
|
||
raise HTTPException(422, "no active intervention to export")
|
||
mode = interventions.mode
|
||
if mode in ("readthrough", "exact"):
|
||
export_fn, kwargs = editing.export_rebase, {"exact": mode == "exact"}
|
||
elif mode == "abliteration":
|
||
export_fn, kwargs = editing.export_abliteration, {}
|
||
else:
|
||
raise HTTPException(422, "export requires a pure-weights mode")
|
||
source_dir = resolve_local_dir(manager.meta["model_id"])
|
||
try:
|
||
await asyncio.to_thread(
|
||
export_fn, rules, manager.jl, manager.meta,
|
||
fmt="full", name=f"{req.name}/hf", source_dir=source_dir,
|
||
scale=interventions.global_scale, **kwargs,
|
||
)
|
||
except ValueError as exc:
|
||
raise HTTPException(422, str(exc))
|
||
baked = "baked"
|
||
|
||
_gguf_state.update(state="running", name=req.name, step="starting", error=None, result=None)
|
||
threading.Thread(
|
||
target=_gguf_worker,
|
||
args=(req.name, req.gguf_type, hf_dir, convert, quantize, gguf_py),
|
||
daemon=True,
|
||
).start()
|
||
return {"started": True, "checkpoint": baked, "state": dict(_gguf_state)}
|
||
|
||
|
||
class GGUFCacheRequest(BaseModel):
|
||
name: str
|
||
|
||
|
||
@app.post("/api/edit/gguf-cache/delete")
|
||
def api_gguf_cache_delete(req: GGUFCacheRequest):
|
||
"""Drop the cached HF checkpoint of a GGUF export (the .gguf files stay)."""
|
||
import shutil
|
||
hf_dir = editing.EDITS_DIR / _safe_name(req.name) / "hf"
|
||
if not hf_dir.is_dir():
|
||
raise HTTPException(404, f"no cached checkpoint for {req.name}")
|
||
if _gguf_state["state"] == "running" and _gguf_state["name"] == req.name:
|
||
raise HTTPException(409, "a GGUF export is using this cache")
|
||
freed = sum(f.stat().st_size for f in hf_dir.rglob("*") if f.is_file())
|
||
shutil.rmtree(hf_dir)
|
||
return {"deleted": str(hf_dir), "freed_bytes": freed}
|
||
|
||
|
||
class GenerateSyncRequest(BaseModel):
|
||
messages: list[dict]
|
||
sampling: dict = {}
|
||
|
||
|
||
@app.post("/api/generate")
|
||
async def api_generate_sync(req: GenerateSyncRequest):
|
||
"""Synchronous generation, no persistence or lens frames: for the CLI tools
|
||
(scripts/jlab.py). Active interventions apply just like in the chat."""
|
||
if manager.hf_model is None:
|
||
raise HTTPException(422, "no model loaded")
|
||
if manager.busy:
|
||
raise HTTPException(409, f"busy: {manager.busy}")
|
||
done = {}
|
||
|
||
def emit(frame):
|
||
if frame["type"] == "done":
|
||
done.update(frame)
|
||
elif frame["type"] == "error":
|
||
done["error"] = frame.get("message")
|
||
|
||
try:
|
||
await asyncio.to_thread(
|
||
manager.generate,
|
||
req.messages,
|
||
req.sampling,
|
||
threading.Event(),
|
||
emit,
|
||
lens=None,
|
||
ablator=interventions if interventions.active else None,
|
||
)
|
||
except Exception as exc:
|
||
raise HTTPException(500, str(exc))
|
||
if done.get("error"):
|
||
raise HTTPException(500, done["error"])
|
||
return {"text": done.get("text", ""), "stats": done.get("stats")}
|
||
|
||
|
||
class NeighborsRequest(BaseModel):
|
||
token_ids: list[int]
|
||
k: int = 3
|
||
|
||
|
||
@app.post("/api/token-neighbors")
|
||
async def api_token_neighbors(req: NeighborsRequest):
|
||
if manager.hf_model is None:
|
||
raise HTTPException(422, "no model loaded")
|
||
key = ((manager.meta or {}).get("model_id"), (manager.meta or {}).get("revision"))
|
||
try:
|
||
result = await asyncio.to_thread(
|
||
neighbors.lookup, manager.jl, manager.tokenizer, key,
|
||
req.token_ids[:64], req.k,
|
||
)
|
||
except ValueError as exc:
|
||
raise HTTPException(422, str(exc))
|
||
return {"neighbors": {str(tid): entries for tid, entries in result.items()}}
|
||
|
||
|
||
@app.get("/api/token-lookup")
|
||
def api_token_lookup(q: str):
|
||
if manager.tokenizer is None:
|
||
raise HTTPException(422, "no model loaded")
|
||
tokenizer = manager.tokenizer
|
||
candidates = {}
|
||
for variant in (q, " " + q, q.lower(), " " + q.lower(),
|
||
q.capitalize(), " " + q.capitalize(), q.upper(), " " + q.upper()):
|
||
ids = tokenizer.encode(variant, add_special_tokens=False)
|
||
if len(ids) == 1 and ids[0] not in candidates:
|
||
candidates[ids[0]] = tokenizer.decode([ids[0]])
|
||
return {"candidates": [{"id": tid, "str": s} for tid, s in candidates.items()]}
|
||
|
||
|
||
@app.get("/api/registry/local")
|
||
def api_registry_local():
|
||
return {"lenses": registry.local_lenses()}
|
||
|
||
|
||
@app.get("/api/registry/for-model")
|
||
async def api_registry_for_model(model_id: str, revision: str | None = None):
|
||
return await asyncio.to_thread(registry.lenses_for_model, model_id, revision)
|
||
|
||
|
||
@app.get("/api/registry/resolve")
|
||
def api_registry_resolve(path: str | None = None, repo_id: str | None = None, filename: str | None = None):
|
||
return registry.resolve_lens(path=path, repo_id=repo_id, filename=filename)
|
||
|
||
|
||
@app.post("/api/fit")
|
||
def api_fit(req: FitRequest):
|
||
if manager.hf_model is not None:
|
||
raise HTTPException(
|
||
409, "unload the model first: fitting needs all the VRAM"
|
||
)
|
||
valid = _valid_devices()
|
||
bad = [d for d in req.devices if d not in valid]
|
||
if bad:
|
||
raise HTTPException(422, f"invalid device(s): {', '.join(bad)}")
|
||
source = resolve_source(req.model_id)
|
||
try:
|
||
return fit_manager.start(
|
||
model_id=req.model_id,
|
||
source=source,
|
||
model_revision=_resolve_revision(source),
|
||
n_prompts=req.n_prompts,
|
||
dtype=req.dtype,
|
||
quant=req.quant,
|
||
datasets=req.datasets,
|
||
devices=req.devices,
|
||
name=req.name,
|
||
dim_batch=req.dim_batch,
|
||
max_seq_len=req.max_seq_len,
|
||
source_layers=req.source_layers,
|
||
continue_from=req.continue_from,
|
||
)
|
||
except ValueError as exc:
|
||
raise HTTPException(409, str(exc))
|
||
|
||
|
||
@app.get("/api/fit/status")
|
||
def api_fit_status():
|
||
return fit_manager.state
|
||
|
||
|
||
@app.post("/api/fit/stop")
|
||
def api_fit_stop():
|
||
return fit_manager.stop()
|
||
|
||
|
||
@app.post("/api/lens/pin")
|
||
async def api_lens_pin(req: PinRequest):
|
||
if manager.busy:
|
||
raise HTTPException(409, f"busy: {manager.busy}")
|
||
try:
|
||
return await asyncio.to_thread(
|
||
lens_manager.pin_ranks, req.gen_id, req.token_ids, manager.jl
|
||
)
|
||
except ValueError as exc:
|
||
raise HTTPException(422, str(exc))
|
||
|
||
|
||
# Alternative/duplicate weight folders, never needed for transformers inference
|
||
# (GPT-OSS-20B ships original/ + metal/ = 2 × ~14 GB of waste).
|
||
DOWNLOAD_IGNORE_DIRS = [
|
||
"original/*", "metal/*", "onnx/*", "openvino/*", "coreml/*", "gguf/*",
|
||
]
|
||
# Auxiliary files that are always useful (configs, tokenizer, custom code) — light.
|
||
DOWNLOAD_EXTRAS = ["*.json", "*.txt", "*.model", "*.tiktoken", "*.jinja", "*.py", "*.md"]
|
||
# "large" fp32 model: past 1 GB of weights, convert to bf16 automatically
|
||
AUTO_BF16_MIN_BYTES = 1_000_000_000
|
||
|
||
|
||
def _plan_download(api, repo_id, token):
|
||
"""Pick the strict minimum: ONE weight set (the lightest if several variants)
|
||
+ the auxiliary files. Returns (allow_patterns, ignore_patterns, plan) —
|
||
allow_patterns None = "take everything" fallback."""
|
||
import re
|
||
|
||
files = {} # path -> size
|
||
for entry in api.list_repo_tree(repo_id, recursive=True, token=token or None):
|
||
size = getattr(entry, "size", None)
|
||
if size is not None:
|
||
files[entry.path] = size
|
||
|
||
def in_ignored_dir(path):
|
||
return any(path.startswith(d.split("/*")[0] + "/") for d in DOWNLOAD_IGNORE_DIRS)
|
||
|
||
# root-level safetensors sets, grouped by variant:
|
||
# "model(-00001-of-00002)?.safetensors" -> group "model";
|
||
# "model.fp32(-...)?.safetensors" -> group "model.fp32", etc.
|
||
st_groups = {}
|
||
for path, size in files.items():
|
||
if "/" in path or not path.endswith(".safetensors"):
|
||
continue
|
||
stem = re.sub(r"-\d{5}-of-\d{5}", "", path.removesuffix(".safetensors"))
|
||
st_groups.setdefault(stem, []).append(path)
|
||
|
||
if st_groups:
|
||
stem, chosen = min(
|
||
st_groups.items(), key=lambda kv: sum(files[p] for p in kv[1])
|
||
)
|
||
index = f"{stem}.safetensors.index.json"
|
||
patterns = sorted(chosen) + ([index] if index in files else []) + DOWNLOAD_EXTRAS
|
||
return patterns, None, {
|
||
"kind": "safetensors",
|
||
"variant": stem,
|
||
"size_bytes": sum(files[p] for p in chosen),
|
||
}
|
||
|
||
# no safetensors (legacy .bin/.h5 repos, or GGUF-only ones we can't load):
|
||
# take everything except the alternative folders and obvious format
|
||
# duplicates. GGUF weights are ignored — J-Wash only loads transformers
|
||
# (safetensors) models.
|
||
ignore = DOWNLOAD_IGNORE_DIRS + ["*.gguf", "*.msgpack", "*.h5", "*.tflite", "*.onnx"]
|
||
return None, ignore, {
|
||
"kind": "fallback",
|
||
"size_bytes": sum(
|
||
s for p, s in files.items()
|
||
if not in_ignored_dir(p) and not p.endswith(".gguf")
|
||
),
|
||
}
|
||
|
||
|
||
def _maybe_autoconvert_bf16(repo_id, state):
|
||
"""After download: if the safetensors weights are float32 and heavy, convert
|
||
to bf16 automatically into a local folder (halves the space in use; the HF
|
||
cache source stays intact)."""
|
||
from pathlib import Path
|
||
|
||
from safetensors import safe_open
|
||
|
||
from core.model_manager import convert_to_bf16, resolve_local_dir
|
||
|
||
src = resolve_local_dir(repo_id)
|
||
if not src:
|
||
return
|
||
src = Path(src)
|
||
shards = sorted(src.glob("*.safetensors"))
|
||
if not shards:
|
||
return
|
||
total = sum(s.stat().st_size for s in shards)
|
||
if total < AUTO_BF16_MIN_BYTES:
|
||
return
|
||
import math
|
||
|
||
with safe_open(str(shards[0]), framework="pt") as f:
|
||
keys = list(f.keys())
|
||
if not keys:
|
||
return
|
||
# the shard's biggest tensor is representative of the "large layers"
|
||
biggest = max(keys, key=lambda k: math.prod(f.get_slice(k).get_shape()))
|
||
dtype = str(f.get_slice(biggest).get_dtype())
|
||
if dtype not in ("F32", "F64"):
|
||
return
|
||
state.update(state="converting")
|
||
base = repo_id.split("/")[-1]
|
||
result = convert_to_bf16(str(src), out_dir=str(config.LOCAL_MODELS_ROOT / f"{base}-bf16"))
|
||
state.update(converted=result["id"])
|
||
|
||
|
||
def _download_worker(repo_id):
|
||
import os
|
||
|
||
from huggingface_hub import HfApi, snapshot_download
|
||
from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError
|
||
|
||
token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
|
||
state = _downloads[repo_id]
|
||
try:
|
||
allow, ignore, plan = None, None, None
|
||
try:
|
||
allow, ignore, plan = _plan_download(HfApi(), repo_id, token)
|
||
state.update(plan=plan)
|
||
except Exception:
|
||
# planning failed (network, permissions): cautious fallback
|
||
ignore = DOWNLOAD_IGNORE_DIRS + ["*.gguf", "*.msgpack", "*.h5", "*.tflite", "*.onnx"]
|
||
# progress: poll the cache size (robust — the tqdm hook misses some files
|
||
# depending on the download mechanism). We sum the repo's blobs (including
|
||
# .incomplete files) and compare to the planned total.
|
||
total_bytes = (plan or {}).get("size_bytes") or 0
|
||
_blobs = config.HF_CACHE / "hub" / f"models--{repo_id.replace('/', '--')}" / "blobs"
|
||
stop_poll = threading.Event()
|
||
|
||
def _poll_progress():
|
||
while not stop_poll.is_set():
|
||
done = 0
|
||
if _blobs.exists():
|
||
for f in _blobs.iterdir():
|
||
try:
|
||
done += f.stat().st_size
|
||
except OSError:
|
||
pass
|
||
if total_bytes:
|
||
state["progress"] = {"done": min(done, total_bytes), "total": total_bytes}
|
||
stop_poll.wait(1.0)
|
||
|
||
poller = threading.Thread(target=_poll_progress, daemon=True)
|
||
poller.start()
|
||
try:
|
||
snapshot_download(
|
||
repo_id, token=token or None,
|
||
allow_patterns=allow, ignore_patterns=ignore,
|
||
)
|
||
finally:
|
||
stop_poll.set()
|
||
state.pop("progress", None)
|
||
if plan and plan["kind"] == "safetensors":
|
||
_maybe_autoconvert_bf16(repo_id, state)
|
||
state.update(state="done", error=None)
|
||
except GatedRepoError:
|
||
msg = (
|
||
f'gated repo "{repo_id}": accept the terms on huggingface.co and make '
|
||
"sure a valid HF_TOKEN is set in the environment."
|
||
+ ("" if token else " (no HF_TOKEN detected)")
|
||
)
|
||
state.update(state="error", error=msg)
|
||
except RepositoryNotFoundError:
|
||
state.update(
|
||
state="error",
|
||
error=f'repo "{repo_id}" not found (or private without access using the current token)',
|
||
)
|
||
except Exception as exc:
|
||
state.update(state="error", error=str(exc))
|
||
|
||
|
||
@app.post("/api/download")
|
||
def api_download(req: DownloadRequest):
|
||
repo_id = req.repo_id.strip()
|
||
with _downloads_lock:
|
||
current = _downloads.get(repo_id)
|
||
if current and current["state"] == "running":
|
||
raise HTTPException(409, f"download already in progress: {repo_id}")
|
||
# several downloads in parallel: one state per repo
|
||
_downloads[repo_id] = {"repo_id": repo_id, "state": "running", "error": None}
|
||
threading.Thread(target=_download_worker, args=(repo_id,), daemon=True).start()
|
||
return _downloads[repo_id]
|
||
|
||
|
||
@app.delete("/api/download/{repo_id:path}")
|
||
def api_download_dismiss(repo_id: str):
|
||
"""Remove a finished (done/error) entry from the displayed list."""
|
||
with _downloads_lock:
|
||
state = _downloads.get(repo_id)
|
||
if state and state["state"] != "running":
|
||
del _downloads[repo_id]
|
||
return {"downloads": list(_downloads.values())}
|
||
|
||
|
||
@app.get("/api/browse")
|
||
def api_browse(path: str | None = None):
|
||
from core.model_manager import browse_dir
|
||
|
||
try:
|
||
return browse_dir(path)
|
||
except ValueError as exc:
|
||
raise HTTPException(404, str(exc))
|
||
|
||
|
||
class PickPathRequest(BaseModel):
|
||
kind: str = "dir" # kept for API compatibility; only directory picking is used
|
||
|
||
|
||
_pick_lock = threading.Lock()
|
||
|
||
|
||
@app.post("/api/pick-path")
|
||
async def api_pick_path(req: PickPathRequest):
|
||
"""Open Windows' NATIVE file picker (the server runs on the user's own
|
||
machine) and return the chosen path — this notably lets you paste a path,
|
||
which the built-in browser cannot do."""
|
||
|
||
def pick():
|
||
if not _pick_lock.acquire(blocking=False):
|
||
raise ValueError("a file picker is already open")
|
||
try:
|
||
# tkinter ships with CPython on every platform, but headless
|
||
# Linux installs may lack it (or a display): fail with a hint
|
||
# instead of a stack trace — the built-in Browse still works.
|
||
try:
|
||
import tkinter as tk
|
||
from tkinter import filedialog
|
||
except ImportError:
|
||
raise ValueError(
|
||
"no native folder picker available (tkinter missing) — "
|
||
"use the built-in Browse, or paste the path directly"
|
||
)
|
||
|
||
root = tk.Tk()
|
||
root.withdraw()
|
||
root.attributes("-topmost", True)
|
||
try:
|
||
path = filedialog.askdirectory(
|
||
parent=root, title="Choose a model folder (HF)"
|
||
)
|
||
finally:
|
||
root.destroy()
|
||
return {"path": path or None}
|
||
finally:
|
||
_pick_lock.release()
|
||
|
||
try:
|
||
return await asyncio.to_thread(pick)
|
||
except ValueError as exc:
|
||
raise HTTPException(409, str(exc))
|
||
|
||
|
||
class ConvertRequest(BaseModel):
|
||
path: str
|
||
|
||
|
||
_convert_state = {"path": None, "state": "idle", "error": None, "result": None}
|
||
|
||
|
||
def _convert_worker(path):
|
||
from core.model_manager import convert_to_bf16
|
||
|
||
try:
|
||
result = convert_to_bf16(path)
|
||
_convert_state.update(state="done", error=None, result=result)
|
||
except Exception as exc:
|
||
_convert_state.update(state="error", error=str(exc))
|
||
|
||
|
||
@app.post("/api/convert-bf16")
|
||
def api_convert_bf16(req: ConvertRequest):
|
||
if _convert_state["state"] == "running":
|
||
raise HTTPException(409, "a conversion is already in progress")
|
||
_convert_state.update(path=req.path, state="running", error=None, result=None)
|
||
threading.Thread(target=_convert_worker, args=(req.path,), daemon=True).start()
|
||
return _convert_state
|
||
|
||
|
||
@app.get("/api/conversations")
|
||
def api_conversations(query: str | None = None):
|
||
return {"conversations": store.list_conversations(query)}
|
||
|
||
|
||
@app.get("/api/conversations/{cid}")
|
||
def api_conversation(cid: int):
|
||
try:
|
||
return store.get_conversation(cid)
|
||
except ValueError as exc:
|
||
raise HTTPException(404, str(exc))
|
||
|
||
|
||
@app.patch("/api/conversations/{cid}")
|
||
def api_conversation_patch(cid: int, req: ConversationPatch):
|
||
store.update_conversation(cid, title=req.title, tags=req.tags)
|
||
return {"ok": True}
|
||
|
||
|
||
@app.delete("/api/conversations/{cid}")
|
||
def api_conversation_delete(cid: int):
|
||
store.delete_conversation(cid)
|
||
return {"ok": True}
|
||
|
||
|
||
@app.get("/api/messages/{mid}/frames")
|
||
def api_message_frames(mid: int):
|
||
try:
|
||
return store.load_frames(mid)
|
||
except ValueError as exc:
|
||
raise HTTPException(404, str(exc))
|
||
|
||
|
||
class MessagePatch(BaseModel):
|
||
content: str
|
||
|
||
|
||
@app.patch("/api/messages/{mid}")
|
||
def api_message_patch(mid: int, req: MessagePatch):
|
||
"""Edit a message's content (e.g. rewrite an assistant reply). Later turns
|
||
are generated from the stored path, so the edit takes effect immediately."""
|
||
try:
|
||
store.update_message(mid, req.content)
|
||
except ValueError as exc:
|
||
raise HTTPException(404, str(exc))
|
||
return {"ok": True, "id": mid}
|
||
|
||
|
||
@app.get("/api/conversations/{cid}/export")
|
||
def api_conversation_export(cid: int, format: str = "json", frames: int = 0):
|
||
try:
|
||
body, media_type = store.export(cid, fmt=format, include_frames=bool(frames))
|
||
except ValueError as exc:
|
||
raise HTTPException(404, str(exc))
|
||
ext = "json" if format == "json" else "md"
|
||
return Response(
|
||
content=body,
|
||
media_type=f"{media_type}; charset=utf-8",
|
||
headers={"Content-Disposition": f'attachment; filename="conversation-{cid}.{ext}"'},
|
||
)
|
||
|
||
|
||
def _generate_safely(messages, sampling, stop_event, emit, lens):
|
||
try:
|
||
manager.generate(
|
||
messages, sampling, stop_event, emit, lens=lens,
|
||
ablator=interventions if interventions.active else None,
|
||
)
|
||
except Exception as exc:
|
||
emit({"type": "error", "message": str(exc)})
|
||
|
||
|
||
def _persisted_generate(req, stop_event, emit, lens):
|
||
try:
|
||
continue_id = req.get("continue_message_id")
|
||
if continue_id is not None:
|
||
_persisted_continue(req, continue_id, stop_event, emit, lens)
|
||
return
|
||
conversation_id = req.get("conversation_id")
|
||
parent_id = req.get("parent_id")
|
||
system = (req.get("system") or "").strip()
|
||
content = req.get("content")
|
||
if conversation_id is None and not content:
|
||
emit({"type": "error", "message": "content required for a new conversation"})
|
||
return
|
||
if conversation_id is None:
|
||
conversation_id = store.create_conversation(content[:60])
|
||
if system:
|
||
parent_id = store.add_message(conversation_id, None, "system", system)
|
||
if content:
|
||
parent_id = store.add_message(conversation_id, parent_id, "user", content)
|
||
if parent_id is None:
|
||
emit({"type": "error", "message": "parent_id or content required"})
|
||
return
|
||
emit({
|
||
"type": "persisted",
|
||
"conversation_id": conversation_id,
|
||
"user_message_id": parent_id,
|
||
})
|
||
context = store.path_to_root(parent_id)
|
||
layers_used = list(lens.layers) if lens is not None else []
|
||
k_used = lens.k if lens is not None else 0
|
||
frames_acc = []
|
||
done_holder = {}
|
||
|
||
def emit_inner(frame):
|
||
if frame["type"] == "done":
|
||
done_holder.update(frame)
|
||
else:
|
||
if frame["type"] == "frame":
|
||
frames_acc.append(frame)
|
||
emit(frame)
|
||
|
||
manager.generate(
|
||
context, req.get("sampling", {}), stop_event, emit_inner, lens=lens,
|
||
ablator=interventions if interventions.active else None,
|
||
)
|
||
meta = dict(
|
||
done_holder.get("meta") or {},
|
||
stats=done_holder.get("stats"),
|
||
stopped=done_holder.get("stopped"),
|
||
)
|
||
message_id = store.add_message(
|
||
conversation_id, parent_id, "assistant", done_holder.get("text", ""), meta=meta
|
||
)
|
||
if frames_acc:
|
||
store.save_frames(message_id, frames_acc, layers_used, k_used)
|
||
emit(dict(done_holder, conversation_id=conversation_id, message_id=message_id))
|
||
except Exception as exc:
|
||
emit({"type": "error", "message": str(exc)})
|
||
|
||
|
||
def _persisted_continue(req, message_id, stop_event, emit, lens):
|
||
"""Extend an existing assistant reply: generate with the turn left open,
|
||
append the text to the message, and merge the new lens frames into its
|
||
stored blob (positions keep increasing, so both parts stay coherent)."""
|
||
msg = store.get_message(message_id)
|
||
if msg["role"] != "assistant":
|
||
emit({"type": "error", "message": "only an assistant reply can be continued"})
|
||
return
|
||
context = store.path_to_root(message_id)
|
||
layers_used = list(lens.layers) if lens is not None else []
|
||
k_used = lens.k if lens is not None else 0
|
||
frames_acc = []
|
||
done_holder = {}
|
||
|
||
def emit_inner(frame):
|
||
if frame["type"] == "done":
|
||
done_holder.update(frame)
|
||
else:
|
||
if frame["type"] == "frame":
|
||
frames_acc.append(frame)
|
||
emit(frame)
|
||
|
||
manager.generate(
|
||
context, req.get("sampling", {}), stop_event, emit_inner, lens=lens,
|
||
ablator=interventions if interventions.active else None,
|
||
continue_final=True,
|
||
)
|
||
new_content = msg["content"] + done_holder.get("text", "")
|
||
meta = json.loads(msg["meta"]) if msg.get("meta") else {}
|
||
meta = dict(
|
||
meta,
|
||
stats=done_holder.get("stats"),
|
||
stopped=done_holder.get("stopped"),
|
||
continued=True,
|
||
)
|
||
store.update_message(message_id, new_content, meta=meta)
|
||
if frames_acc:
|
||
merged = frames_acc
|
||
if msg.get("frames_file"):
|
||
try:
|
||
merged = store.load_frames(message_id)["frames"] + frames_acc
|
||
except Exception:
|
||
pass
|
||
store.save_frames(message_id, merged, layers_used, k_used)
|
||
emit(dict(
|
||
done_holder,
|
||
conversation_id=msg["conversation_id"],
|
||
message_id=message_id,
|
||
text=new_content,
|
||
continued=True,
|
||
))
|
||
|
||
|
||
async def _watch_stop(ws, stop_event):
|
||
while True:
|
||
msg = json.loads(await ws.receive_text())
|
||
if msg.get("type") == "stop":
|
||
stop_event.set()
|
||
|
||
|
||
async def _run_chat(ws, req):
|
||
loop = asyncio.get_running_loop()
|
||
queue = asyncio.Queue()
|
||
stop_event = threading.Event()
|
||
|
||
def emit(frame):
|
||
loop.call_soon_threadsafe(queue.put_nowait, frame)
|
||
|
||
lens = lens_manager if req.get("lens") and lens_manager.lens is not None else None
|
||
if "messages" in req:
|
||
worker = asyncio.create_task(
|
||
asyncio.to_thread(
|
||
_generate_safely, req["messages"], req.get("sampling", {}), stop_event, emit, lens
|
||
)
|
||
)
|
||
else:
|
||
worker = asyncio.create_task(
|
||
asyncio.to_thread(_persisted_generate, req, stop_event, emit, lens)
|
||
)
|
||
receiver = asyncio.create_task(_watch_stop(ws, stop_event))
|
||
try:
|
||
while True:
|
||
frame = await queue.get()
|
||
await _ws_send(ws, json.dumps(frame))
|
||
if frame["type"] in ("done", "error"):
|
||
break
|
||
finally:
|
||
stop_event.set()
|
||
receiver.cancel()
|
||
await asyncio.gather(worker, receiver, return_exceptions=True)
|
||
|
||
|
||
@app.websocket("/ws")
|
||
async def ws_endpoint(ws: WebSocket):
|
||
await ws.accept()
|
||
_ws_locks[ws] = asyncio.Lock()
|
||
try:
|
||
while True:
|
||
req = json.loads(await ws.receive_text())
|
||
if req.get("type") != "chat":
|
||
continue
|
||
if manager.hf_model is None:
|
||
await _ws_send(
|
||
ws, json.dumps({"type": "error", "message": "no model loaded"})
|
||
)
|
||
continue
|
||
if manager.busy:
|
||
await _ws_send(
|
||
ws, json.dumps({"type": "error", "message": f"busy: {manager.busy}"})
|
||
)
|
||
continue
|
||
await _run_chat(ws, req)
|
||
except (WebSocketDisconnect, RuntimeError):
|
||
pass
|
||
finally:
|
||
_ws_locks.pop(ws, None)
|
||
|
||
|
||
if config.UI_DIST.exists():
|
||
app.mount("/", StaticFiles(directory=config.UI_DIST, html=True), name="ui")
|