Joey Grasty
e25ee046a3
llama.cpp-based implementation of the countbayesie Future-Entropy Sampler, built via a direct ctypes binding to libllama.so. Includes a one-shot CLI (entropy_cli.py), an OpenAI-compatible server (entropy_server.py), and the tuning/benchmark scripts used to derive the alpha/confidence-threshold defaults documented in the README. Model path and llama.cpp library path are now read from environment variables (ENTROPY_SAMPLER_MODEL, LLAMA_CPP_LIB) instead of being hardcoded, and gguf-py is pulled from PyPI instead of a local llama.cpp checkout, so this runs on any machine with a compatible llama.cpp build and model.
326 lines
14 KiB
Python
326 lines
14 KiB
Python
"""OpenAI-compatible HTTP server for the Future-Entropy Sampler prototype -
|
|
meant to run alongside the stock qwen3.5-llama-start server (port 30000) so
|
|
the two can be compared side by side from the same chat UI (e.g. Open WebUI
|
|
as two separate model connections).
|
|
|
|
Sampler configuration (alpha/sine, top_k, confidence_threshold, ...) is fixed
|
|
at startup via CLI flags, same spirit as entropy_cli.py - there is no
|
|
per-request override and no runtime reconfiguration; restart with different
|
|
flags to try a different configuration.
|
|
|
|
Usage:
|
|
python entropy_server.py --port 30001
|
|
python entropy_server.py --port 30001 --sine 16,0.6,-0.2 --confidence-threshold 0.8
|
|
"""
|
|
import argparse
|
|
import json
|
|
import sys
|
|
import threading
|
|
import time
|
|
import uuid
|
|
|
|
import gguf
|
|
import jinja2
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.responses import StreamingResponse
|
|
from pydantic import BaseModel
|
|
import uvicorn
|
|
|
|
import llama_capi as C
|
|
from entropy_sampler import EntropySampler, sine_alpha_schedule, MODEL_PATH
|
|
from entropy_cli import parse_sine, SAFE_PEAK_ALPHA
|
|
|
|
|
|
DEFAULT_MODEL_ID = "qwen3.5-35b-a3b-entropy"
|
|
|
|
|
|
def load_chat_template(model_path):
|
|
reader = gguf.GGUFReader(model_path)
|
|
field = reader.fields.get("tokenizer.chat_template")
|
|
if field is None:
|
|
raise RuntimeError(f"{model_path} has no embedded tokenizer.chat_template")
|
|
return bytes(field.parts[-1]).decode("utf-8")
|
|
|
|
|
|
def _raise_exception(msg):
|
|
# The Qwen chat template calls raise_exception(...) as a Jinja global,
|
|
# matching the helper HF's apply_chat_template injects into its env.
|
|
raise ValueError(msg)
|
|
|
|
|
|
def build_jinja_template(template_src):
|
|
env = jinja2.Environment()
|
|
env.globals["raise_exception"] = _raise_exception
|
|
env.filters["tojson"] = lambda v: json.dumps(v)
|
|
return env.from_string(template_src)
|
|
|
|
|
|
class ChatMessage(BaseModel):
|
|
role: str
|
|
content: str = ""
|
|
|
|
|
|
class ChatCompletionRequest(BaseModel):
|
|
model: str | None = None
|
|
messages: list[ChatMessage]
|
|
stream: bool = False
|
|
max_tokens: int | None = None
|
|
# Accepted for OpenAI-client compatibility, intentionally ignored: this
|
|
# sampler's behavior is fully determined by the server's baked-in
|
|
# alpha/sine/top_k/confidence_threshold config, not per-request params.
|
|
temperature: float | None = None
|
|
top_p: float | None = None
|
|
|
|
|
|
def build_arg_parser():
|
|
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
p.add_argument("--host", default="0.0.0.0")
|
|
p.add_argument("--port", type=int, default=30001)
|
|
p.add_argument("--model", default=MODEL_PATH)
|
|
p.add_argument("--model-id", default=DEFAULT_MODEL_ID,
|
|
help="Model name reported via /v1/models and in responses.")
|
|
|
|
alpha_group = p.add_mutually_exclusive_group()
|
|
alpha_group.add_argument("--alpha", type=float, default=None,
|
|
help="Fixed alpha in [-1, 1]. Overrides --sine if given.")
|
|
alpha_group.add_argument("--sine", type=parse_sine, default="16,0.6,-0.2",
|
|
metavar="PERIOD,AMPLITUDE[,OFFSET[,PHASE]]",
|
|
help="Default: 16,0.6,-0.2 (peak +0.4) - the best-tested "
|
|
"config this session. Keep peak (offset+amplitude) "
|
|
"<= ~+0.5 for safety.")
|
|
|
|
p.add_argument("--top-k", type=int, default=12)
|
|
p.add_argument("--top-n-future", type=int, default=20)
|
|
p.add_argument("--confidence-threshold", type=float, default=0.8)
|
|
p.add_argument("--n-ctx", type=int, default=212992,
|
|
help="Requested total context, split n_seq_max=(top_k+1) ways since "
|
|
"this model's KV cache is not unified across fork sequences "
|
|
"(kv_unified=false). Default 212992 -> ~16384 tokens/sequence "
|
|
"at top_k=12. Since each chat request re-sends the full "
|
|
"conversation history and re-primes from scratch, this is the "
|
|
"budget for prompt+reply *per request*, and it must also cover "
|
|
"the whole growing history in a multi-turn conversation.")
|
|
p.add_argument("--n-batch", type=int, default=16384,
|
|
help="Hard ceiling on tokens submitted in one llama_decode() call. "
|
|
"prime() decodes the whole rendered prompt/history in a single "
|
|
"call, so this must cover the largest expected prompt. Default "
|
|
"16384 matches the default per-sequence context budget. Too "
|
|
"small and long prompts trigger GGML_ASSERT(n_tokens_all <= "
|
|
"n_batch), which aborts the whole process (not a catchable "
|
|
"Python exception) - this happened with the llama.cpp default "
|
|
"of 2048 on a 2059-token prompt.")
|
|
p.add_argument("--n-ubatch", type=int, default=512,
|
|
help="Internal per-step chunk size within a decode call - llama.cpp "
|
|
"processes n_batch in n_ubatch-sized pieces automatically. "
|
|
"Smaller = less peak compute-buffer memory.")
|
|
p.add_argument("--n-threads", type=int, default=8)
|
|
p.add_argument("--max-tokens-default", type=int, default=3000,
|
|
help="Used when a request doesn't specify max_tokens (Open WebUI "
|
|
"typically doesn't). Sized for ~1600-word creative-writing "
|
|
"replies (~2000-2200 tokens) with margin - 512 and 1536 both "
|
|
"proved too tight and truncated mid-sentence.")
|
|
p.add_argument("--think-ban", dest="think_ban", action="store_true", default=True)
|
|
p.add_argument("--no-think-ban", dest="think_ban", action="store_false")
|
|
p.add_argument("--enable-thinking", dest="enable_thinking", action="store_true", default=False,
|
|
help="Chat template's enable_thinking flag. Default off - this "
|
|
"sampler is tuned for direct creative-writing continuation, "
|
|
"not reasoning mode.")
|
|
return p
|
|
|
|
|
|
def make_alpha(args):
|
|
if args.alpha is not None:
|
|
if args.alpha >= 0.7:
|
|
print(f"warning: fixed alpha {args.alpha:+.2f} >= +0.7 reliably degenerated "
|
|
f"into register-breaks in prior testing", file=sys.stderr)
|
|
return lambda: args.alpha, f"fixed={args.alpha:+.2f}"
|
|
|
|
s = args.sine
|
|
peak = s["offset"] + s["amplitude"]
|
|
if peak > SAFE_PEAK_ALPHA:
|
|
print(f"warning: peak alpha {peak:+.2f} exceeds the ~+{SAFE_PEAK_ALPHA:.1f} "
|
|
f"safety margin observed for coherence - expect elevated crash risk",
|
|
file=sys.stderr)
|
|
desc = (f"sine(period={s['period_tokens']:g}, amplitude={s['amplitude']:.2f}, "
|
|
f"offset={s['offset']:.2f}, phase={s['phase']:.2f}) peak={peak:+.2f}")
|
|
return (lambda: sine_alpha_schedule(**s)), desc
|
|
|
|
|
|
def main():
|
|
args = build_arg_parser().parse_args()
|
|
alpha_factory, alpha_desc = make_alpha(args)
|
|
|
|
print("=== entropy server ===")
|
|
print(f"model: {args.model}")
|
|
print(f"model_id: {args.model_id}")
|
|
print(f"top_k: {args.top_k} (n_seq_max={args.top_k + 1})")
|
|
print(f"top_n_future: {args.top_n_future}")
|
|
print(f"alpha: {alpha_desc}")
|
|
print(f"confidence_thr: {args.confidence_threshold}")
|
|
print(f"think_ban: {args.think_ban}")
|
|
print(f"enable_thinking: {args.enable_thinking} (chat template flag)")
|
|
print(f"n_ctx: requested={args.n_ctx}")
|
|
print(f"n_batch: {args.n_batch} n_ubatch: {args.n_ubatch}")
|
|
|
|
template = build_jinja_template(load_chat_template(args.model))
|
|
|
|
sampler = EntropySampler(
|
|
model_path=args.model,
|
|
n_ctx=args.n_ctx,
|
|
top_k=args.top_k,
|
|
top_n_future=args.top_n_future,
|
|
n_threads=args.n_threads,
|
|
think_ban=args.think_ban,
|
|
n_batch=args.n_batch,
|
|
n_ubatch=args.n_ubatch,
|
|
)
|
|
n_ctx_actual = C.lib.llama_n_ctx(sampler.ctx)
|
|
n_ctx_seq = C.lib.llama_n_ctx_seq(sampler.ctx)
|
|
print(f"n_ctx: actual={n_ctx_actual} per-seq={n_ctx_seq}")
|
|
print(f"think_ban ids: {sampler.banned_token_ids}")
|
|
print("=" * 40)
|
|
|
|
# Each request re-primes seq 0 from scratch with the *entire* rendered chat
|
|
# history, so this per-sequence budget has to cover a full growing
|
|
# conversation, not just one reply. Undersizing this is exactly what
|
|
# produced the mid-generation "llama_decode failed rc=1 / failed to find
|
|
# a memory slot" crash this was added after.
|
|
if n_ctx_seq < args.max_tokens_default + 512:
|
|
print(
|
|
f"warning: per-sequence budget ({n_ctx_seq}) is only "
|
|
f"{n_ctx_seq - args.max_tokens_default} tokens above max_tokens_default "
|
|
f"({args.max_tokens_default}) - long prompts or multi-turn history will "
|
|
f"exhaust it. Raise --n-ctx.",
|
|
file=sys.stderr,
|
|
)
|
|
if args.n_batch < n_ctx_seq:
|
|
print(
|
|
f"warning: --n-batch ({args.n_batch}) is smaller than the per-sequence "
|
|
f"context budget ({n_ctx_seq}) - a prompt/history longer than --n-batch "
|
|
f"tokens will hit GGML_ASSERT(n_tokens_all <= n_batch) and abort the whole "
|
|
f"process, even though it would otherwise fit in context. Raise --n-batch "
|
|
f"to at least {n_ctx_seq}.",
|
|
file=sys.stderr,
|
|
)
|
|
|
|
gen_lock = threading.Lock()
|
|
|
|
def render_prompt(messages):
|
|
try:
|
|
return template.render(
|
|
messages=messages, tools=None, add_generation_prompt=True,
|
|
enable_thinking=args.enable_thinking,
|
|
)
|
|
except ValueError as e:
|
|
raise HTTPException(status_code=400, detail=str(e))
|
|
|
|
app = FastAPI()
|
|
|
|
@app.get("/v1/models")
|
|
def list_models():
|
|
return {
|
|
"object": "list",
|
|
"data": [{"id": args.model_id, "object": "model", "owned_by": "local"}],
|
|
}
|
|
|
|
@app.post("/v1/chat/completions")
|
|
def chat_completions(req: ChatCompletionRequest):
|
|
messages = [m.model_dump() for m in req.messages]
|
|
prompt = render_prompt(messages)
|
|
max_new_tokens = req.max_tokens or args.max_tokens_default
|
|
|
|
completion_id = f"chatcmpl-{uuid.uuid4().hex}"
|
|
created = int(time.time())
|
|
|
|
n_prompt_tokens = len(C.tokenize(sampler.vocab, prompt))
|
|
print(
|
|
f"[{completion_id}] request: prompt_tokens={n_prompt_tokens} "
|
|
f"requested_max_tokens={req.max_tokens} resolved_max_new_tokens={max_new_tokens} "
|
|
f"stream={req.stream}"
|
|
)
|
|
|
|
if req.stream:
|
|
def event_stream():
|
|
first = {
|
|
"id": completion_id, "object": "chat.completion.chunk", "created": created,
|
|
"model": args.model_id,
|
|
"choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
|
|
}
|
|
yield f"data: {json.dumps(first)}\n\n"
|
|
|
|
n_yielded = 0
|
|
out_of_context = False
|
|
with gen_lock:
|
|
try:
|
|
for piece, _diag in sampler.generate_stream(
|
|
prompt, max_new_tokens=max_new_tokens,
|
|
alpha=alpha_factory(), confidence_threshold=args.confidence_threshold,
|
|
):
|
|
n_yielded += 1
|
|
chunk = {
|
|
"id": completion_id, "object": "chat.completion.chunk", "created": created,
|
|
"model": args.model_id,
|
|
"choices": [{"index": 0, "delta": {"content": piece}, "finish_reason": None}],
|
|
}
|
|
yield f"data: {json.dumps(chunk)}\n\n"
|
|
except RuntimeError as e:
|
|
out_of_context = True
|
|
print(f"warning: generation truncated after {n_yielded} tokens: {e} "
|
|
f"(per-sequence budget exhausted - raise --n-ctx)", file=sys.stderr)
|
|
|
|
finish_reason = "length" if (out_of_context or n_yielded >= max_new_tokens) else "stop"
|
|
print(f"[{completion_id}] done: tokens_generated={n_yielded} finish_reason={finish_reason} "
|
|
f"out_of_context={out_of_context}")
|
|
final = {
|
|
"id": completion_id, "object": "chat.completion.chunk", "created": created,
|
|
"model": args.model_id,
|
|
"choices": [{"index": 0, "delta": {}, "finish_reason": finish_reason}],
|
|
}
|
|
yield f"data: {json.dumps(final)}\n\n"
|
|
yield "data: [DONE]\n\n"
|
|
|
|
return StreamingResponse(event_stream(), media_type="text/event-stream")
|
|
|
|
out_of_context = False
|
|
with gen_lock:
|
|
pieces = []
|
|
try:
|
|
for piece, _diag in sampler.generate_stream(
|
|
prompt, max_new_tokens=max_new_tokens,
|
|
alpha=alpha_factory(), confidence_threshold=args.confidence_threshold,
|
|
):
|
|
pieces.append(piece)
|
|
except RuntimeError as e:
|
|
out_of_context = True
|
|
print(f"warning: generation truncated after {len(pieces)} tokens: {e} "
|
|
f"(per-sequence budget exhausted - raise --n-ctx)", file=sys.stderr)
|
|
text = "".join(pieces)
|
|
finish_reason = "length" if (out_of_context or len(pieces) >= max_new_tokens) else "stop"
|
|
print(f"[{completion_id}] done: tokens_generated={len(pieces)} finish_reason={finish_reason} "
|
|
f"out_of_context={out_of_context}")
|
|
|
|
return {
|
|
"id": completion_id,
|
|
"object": "chat.completion",
|
|
"created": created,
|
|
"model": args.model_id,
|
|
"choices": [{
|
|
"index": 0,
|
|
"message": {"role": "assistant", "content": text},
|
|
"finish_reason": finish_reason,
|
|
}],
|
|
"usage": {
|
|
"prompt_tokens": 0, "completion_tokens": len(pieces),
|
|
"total_tokens": len(pieces),
|
|
},
|
|
}
|
|
|
|
try:
|
|
uvicorn.run(app, host=args.host, port=args.port)
|
|
finally:
|
|
sampler.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|