diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 0baf13fff580..df518ea3eec1 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2314,6 +2314,18 @@ struct server_context_impl { } void send_final_response(server_slot & slot) { + if (slot.task->params.learn_only) { + // a /learn that served its prompt first: the learn outcome is its whole reply + auto res = std::make_unique(); + res->id = slot.task->id; + if (slot.learned) { + res->outcome = *slot.learned; + } else { + res->outcome.failure = "the served prompt finished without its learn step"; + } + queue_results.send(std::move(res)); + return; + } auto res = std::make_unique(); res->id = slot.task->id; @@ -6257,6 +6269,33 @@ void server_routes::init_routes() { return res; } task.learn.reply.assign(tokens.begin() + prefix, tokens.end()); + // A past example (S4 (f)): the slot may no longer hold this prompt. Serve it first + // through the ordinary completion path (prompt cache, checkpoints, the slot's own + // admission), one token, and learn as it finishes: when the slot still holds it, the + // cache makes that free; when it does not, it is one prefill. One prefill path, one + // learn path. The slot is the caller's choice: whose cached prompt it replaces. + server_task serve(SERVER_TASK_TYPE_COMPLETION); + serve.tokens = server_tokens(task.learn.prompt, false); + serve.params = server_schema::eval_llama_cmpl_schema(ctx_server.vocab, params, meta->logit_bias_eog, + json { { "n_predict", 1 }, { "temperature", 0.0 }, { "cache_prompt", true } }); + serve.params.learn = task.learn; + serve.params.learn_only = true; + serve.params.res_type = TASK_RESPONSE_TYPE_NONE; + serve.id_slot = task.learn.id_slot; + auto & rd = res->rd; + serve.id = rd.get_new_id(); + rd.post_task(std::move(serve)); + auto result = rd.next(req.should_stop); + if (!result) { + GGML_ASSERT(req.should_stop()); + return res; + } + if (result->is_error()) { + res->error(result->to_json()); + return res; + } + res->ok(result->to_json()); + return res; } else if (body.contains("reply_tokens")) { task.learn.reply = body.at("reply_tokens").get>(); } else { diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 40107fdf7c99..0defcbd116d5 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -183,6 +183,8 @@ struct task_params { // learn a reply on this request's slot when it finishes, before the slot is released (no other // request can take the slot in between): POST /learn's work, atomic with the turn std::optional learn; + // this completion exists only to serve its prompt for `learn`: its reply is the learn outcome + bool learn_only = false; }; // struct for tracking the state of a task (e.g., for streaming) diff --git a/tools/server/tests/learn-every-format.py b/tools/server/tests/learn-every-format.py index 17721774c87f..8677aed0c963 100755 --- a/tools/server/tests/learn-every-format.py +++ b/tools/server/tests/learn-every-format.py @@ -116,6 +116,25 @@ def landed(lora, outcome): "next": nxt["step"], "stale_generation": stale_gen["failure"], "resumed": resumed["step"]})) sys.exit(0) +# A past example (S4 (f); needs --jinja): learn a conversation the slot no longer holds. The engine +# serves its prompt first through the ordinary completion path, then steps; a second learn of a +# conversation the slot now holds reuses the cached prompt. +if "--past-example" in sys.argv: + raw("/completion", {"prompt": "Something unrelated entirely: the weather in Lisbon", "id_slot": 0, "n_predict": 1, "temperature": 0}) + def lesson(lora, messages): + step = fenced(lora) + return {"messages": messages, "id_slot": 0, "lora_id": lora, "lr": step["lr"], "execution": step["execution"], + "base_step": step["base_step"], "engine": ENGINE, "input": step["input"], + "chat_template_kwargs": {"enable_thinking": False}} + past = [{"role": "user", "content": "Which city is the capital of Kansas?"}, + {"role": "assistant", "content": "Topeka is the capital of Kansas."}] + first = landed(sid, json.loads(raw("/learn", lesson(sid, past)))) + assert first["success"] and first["n_labels"] > 1 and first["step"] == 1, first + again = landed(sid, json.loads(raw("/learn", lesson(sid, past)))) + assert again["success"] and again["step"] == 2 and again["loss"] < first["loss"], (first, again) + print(json.dumps({"past_example": first, "again_cached": again})) + sys.exit(0) + # Two minds on one serving context: each shadow keeps its own optimizer while the other learns in # between (needs shadow-05.gguf and a copy of it, shadow-06.gguf, in --train-dir). if "--two-minds" in sys.argv: