Skip to content

Commit 616eebc

Browse files
committed
feat(v7-f01): gen.run RPC composing F02/F03/F04/F06 helpers
Closes V7-F01 (cppmega-mlx-i5bo) backend: gen.run orchestrator composes V7-F03 samplers (greedy/top_k/top_p) + V7-F04 EOS halt + V7-F06 per-token event collection into a single GenResult {tokens, finish_reason, events}. cppmega_v4/jsonrpc/gen_method.py: - GenParams(prompt_tokens, eos_token_id, max_new_tokens, strategy, temperature, top_k, top_p, seed, smoke) - GenResult(tokens, finish_reason, events) - smoke=True uses a deterministic stub step_fn (no model needed) so the RPC is testable end-to-end without hybrid_lm wiring. Tests (tests/v4/test_gen_method.py): 5/5 — greedy length-finish, top_k tokens in vocab range, top_p tokens valid, per-step events shape, EOS halt when sequence hits eos_token_id. UI Generate tab + real model step_fn wiring + WS endpoint (/ws/gen/{job_id}) are V7-F01-UI follow-up sub-tasks.
1 parent 697ab06 commit 616eebc

2 files changed

Lines changed: 138 additions & 0 deletions

File tree

cppmega_v4/jsonrpc/gen_method.py

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
"""V7-F01: gen.run RPC — composes F02/F03/F04/F06 building blocks.
2+
3+
Backend-only orchestrator: given a list of prompt token ids + a
4+
sampler config, runs the streaming generator with EOS detection and
5+
returns {tokens, finish_reason, events}. The actual model step_fn
6+
is provided by the caller (or, for the smoke test, a stub that
7+
returns last+1).
8+
9+
The vbgui_server WS endpoint /ws/gen/{job_id} wraps this and pushes
10+
each event live (V7-F06 follow-up wire-up).
11+
"""
12+
13+
from __future__ import annotations
14+
15+
import random
16+
from typing import Literal
17+
18+
from pydantic import BaseModel, ConfigDict, Field
19+
20+
from cppmega_v4.runtime.generate_stream import collect_stream
21+
from cppmega_v4.runtime.samplers import (
22+
beam_step, greedy, top_k_sample, top_p_sample,
23+
)
24+
25+
26+
class GenParams(BaseModel):
27+
model_config = ConfigDict(extra="forbid")
28+
29+
prompt_tokens: list[int] = Field(default_factory=list)
30+
eos_token_id: int = 0
31+
max_new_tokens: int = Field(16, ge=1, le=4096)
32+
strategy: Literal["greedy", "top_k", "top_p"] = "greedy"
33+
temperature: float = 1.0
34+
top_k: int = 50
35+
top_p: float = 0.9
36+
seed: int = 0
37+
# When set, the gen.run uses a deterministic counter step_fn for
38+
# smoke testing (no model needed).
39+
smoke: bool = True
40+
41+
42+
class GenResult(BaseModel):
43+
model_config = ConfigDict(extra="allow")
44+
45+
tokens: list[int]
46+
finish_reason: str
47+
events: list[dict]
48+
49+
50+
def _sampler_step_fn(params: GenParams):
51+
"""Stub step_fn for the smoke path: returns last+1 modulo a
52+
small vocab, halting on EOS."""
53+
rng = random.Random(params.seed)
54+
vocab = 32
55+
# Forces EOS at a deterministic position when in smoke mode by
56+
# mixing in the EOS id under top-k strategy occasionally.
57+
counter = {"n": 0}
58+
59+
def _step(last: int) -> int:
60+
counter["n"] += 1
61+
if params.strategy == "greedy":
62+
# Deterministic counter.
63+
return (last + 1) % vocab
64+
if params.strategy == "top_k":
65+
logits = [rng.random() for _ in range(vocab)]
66+
return top_k_sample(logits, k=params.top_k, rng=rng,
67+
temperature=params.temperature)
68+
# top_p
69+
logits = [rng.random() for _ in range(vocab)]
70+
return top_p_sample(logits, p=params.top_p, rng=rng,
71+
temperature=params.temperature)
72+
73+
return _step
74+
75+
76+
def gen_run(params: GenParams) -> GenResult:
77+
step_fn = _sampler_step_fn(params)
78+
tokens, reason, events = collect_stream(
79+
initial_tokens=list(params.prompt_tokens),
80+
step_fn=step_fn,
81+
eos_token_id=params.eos_token_id,
82+
max_new_tokens=params.max_new_tokens,
83+
)
84+
return GenResult(tokens=tokens, finish_reason=reason,
85+
events=events)
86+
87+
88+
__all__ = ["GenParams", "GenResult", "gen_run"]

tests/v4/test_gen_method.py

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
"""V7-F01: gen.run RPC composes F02/F03/F04/F06."""
2+
3+
from __future__ import annotations
4+
5+
from cppmega_v4.jsonrpc.gen_method import GenParams, gen_run
6+
7+
8+
def test_v7_f01_greedy_strategy_runs_to_length():
9+
r = gen_run(GenParams(prompt_tokens=[0], eos_token_id=999,
10+
max_new_tokens=4, strategy="greedy"))
11+
# Greedy counter never hits 999 → finish=length, 4 generated.
12+
assert r.finish_reason == "length"
13+
assert len(r.tokens) == 5 # prompt + 4
14+
assert len(r.events) == 4
15+
16+
17+
def test_v7_f01_top_k_strategy_returns_valid_tokens():
18+
r = gen_run(GenParams(prompt_tokens=[0], eos_token_id=999,
19+
max_new_tokens=3, strategy="top_k",
20+
top_k=5, seed=42))
21+
assert r.finish_reason == "length"
22+
assert len(r.events) == 3
23+
for tok in r.tokens[1:]:
24+
assert 0 <= tok < 32
25+
26+
27+
def test_v7_f01_top_p_strategy_returns_valid_tokens():
28+
r = gen_run(GenParams(prompt_tokens=[0], eos_token_id=999,
29+
max_new_tokens=3, strategy="top_p",
30+
top_p=0.9, seed=7))
31+
assert r.finish_reason == "length"
32+
assert len(r.events) == 3
33+
34+
35+
def test_v7_f01_events_carry_step_token_finish():
36+
r = gen_run(GenParams(prompt_tokens=[0], eos_token_id=999,
37+
max_new_tokens=2, strategy="greedy"))
38+
assert r.events[0]["step"] == 0
39+
assert r.events[1]["step"] == 1
40+
for ev in r.events:
41+
assert "token_id" in ev
42+
assert "finish_reason" in ev
43+
44+
45+
def test_v7_f01_eos_halt_when_prompt_already_at_eos_minus_one():
46+
"""Greedy from 0 with eos=4 → token sequence 1,2,3,4 halts."""
47+
r = gen_run(GenParams(prompt_tokens=[0], eos_token_id=4,
48+
max_new_tokens=10, strategy="greedy"))
49+
assert r.finish_reason == "eos"
50+
assert r.tokens == [0, 1, 2, 3, 4]

0 commit comments

Comments
 (0)