tacit/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2026-07-06 21:09:15 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge pull request #50 from manmohan659/fix/tool-decode-text-match

Browse Source 
fix(serve): decode-tail text match for tool markers
This commit is contained in:
Manmohan 2026-04-22 17:48:58 -04:00 committed by GitHub
commit f41da418ab
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 57 additions and 66 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

123
modal/serve.py
View File

@ -253,28 +253,21 @@ class Inference:
if len(tokens) > max_context: if len(tokens) > max_context:
tokens = tokens[-max_context:] tokens = tokens[-max_context:]
# Ordinary-text token-id sequences for the tool markers.
# The SFT loader tokenizes assistant content with .encode() (not .encode_special()), # The SFT loader tokenizes assistant content with .encode() (not .encode_special()),
# so the model emits these as multi-token sequences, not single special-token ids. # so these markers are emitted as multi-token byte sequences, and BPE has
# Match on the id sequence directly — more reliable than text (BPE partial UTF-8 # multiple valid tokenizations of the same string — so matching on a single
# can make single-token decodes return empty strings). # expected id sequence is unreliable. Instead we decode the tail of the
tool_start_ids = tuple(self.tokenizer.encode("<|python_start|>")) # generated token stream and search for the marker TEXT.
tool_end_ids = tuple(self.tokenizer.encode("<|python_end|>")) tool_start_str = "<|python_start|>"
tool_end_str = "<|python_end|>"
out_start_str = "<|output_start|>" out_start_str = "<|output_start|>"
out_end_str = "<|output_end|>" out_end_str = "<|output_end|>"
# Suppression: model training has many convs that emit a fake
# <|output_start|>…<|output_end|> after our injected one. We stop the
# turn if we detect another full <|output_start|> sequence emitted
# after our injection.
out_start_ids = tuple(self.tokenizer.encode(out_start_str))
async def stream(): async def stream():
input_ids = torch.tensor([tokens], dtype=torch.long, device=self.device) input_ids = torch.tensor([tokens], dtype=torch.long, device=self.device)
gen_ids: list[int] = [] # everything the MODEL sampled this turn gen_ids: list[int] = [] # everything the MODEL sampled this turn
tool_start_idx = -1 # position in gen_ids where <|python_start|> begins
tool_injected = False # once True, stop detecting further tool calls tool_injected = False # once True, stop detecting further tool calls
injection_end_pos = -1 # index in gen_ids where our injected tokens end pre_injection_len = 0 # len(gen_ids) right before we start injection
def _append_token(tid): def _append_token(tid):
nonlocal input_ids nonlocal input_ids
@ -283,19 +276,13 @@ class Inference:
if input_ids.size(1) > self.config.sequence_len: if input_ids.size(1) > self.config.sequence_len:
input_ids = input_ids[:, -self.config.sequence_len:] input_ids = input_ids[:, -self.config.sequence_len:]
def _match_at(seq: list[int], pos: int, pat: tuple) -> bool: def _decode_tail_text(last_n: int = 40) -> str:
if pos < 0 or pos + len(pat) > len(seq): if not gen_ids:
return False return ""
return tuple(seq[pos:pos + len(pat)]) == pat try:
return self.tokenizer.decode(gen_ids[-last_n:])
def _find_subseq(seq: list[int], pat: tuple, start: int = 0) -> int: except Exception:
L = len(pat) return ""
if L == 0 or len(seq) < start + L:
return -1
for i in range(start, len(seq) - L + 1):
if tuple(seq[i:i + L]) == pat:
return i
return -1
with torch.no_grad(): with torch.no_grad():
num_generated = 0 num_generated = 0
@ -319,7 +306,7 @@ class Inference:
gen_ids.append(token_id) gen_ids.append(token_id)
num_generated += 1 num_generated += 1
# Stream raw decoded text (may be empty for partial BPE bytes — that's OK) # Stream raw decoded text (may be empty for partial BPE bytes)
try: try:
token_text = self.tokenizer.decode([token_id]) token_text = self.tokenizer.decode([token_id])
except Exception: except Exception:
@ -327,46 +314,50 @@ class Inference:
if token_text: if token_text:
yield "data: " + json.dumps({"token": token_text, "gpu": 0}) + "\n\n" yield "data: " + json.dumps({"token": token_text, "gpu": 0}) + "\n\n"
# --- tool-call detection (id-sequence match) --- # --- tool-call detection on decoded-tail text ---
if not tool_injected and tool_start_idx < 0: if not tool_injected:
idx = _find_subseq(gen_ids, tool_start_ids, max(0, len(gen_ids) - len(tool_start_ids) - 2)) # Decode the whole turn-so-far and look for markers
if idx >= 0: try:
tool_start_idx = idx full_text = self.tokenizer.decode(gen_ids)
if not tool_injected and tool_start_idx >= 0: except Exception:
# look for <|python_end|> after the payload full_text = ""
end_idx = _find_subseq(gen_ids, tool_end_ids, tool_start_idx + len(tool_start_ids)) if full_text:
if end_idx >= 0: ps = full_text.rfind(tool_start_str)
payload_ids = gen_ids[tool_start_idx + len(tool_start_ids):end_idx] if ps >= 0:
try: pe = full_text.find(tool_end_str, ps + len(tool_start_str))
payload_text = self.tokenizer.decode(payload_ids) if pe >= 0:
invocation = self._parse_tool_call(payload_text) payload_text = full_text[ps + len(tool_start_str):pe]
result = self.tool_registry.execute(invocation.tool_name, invocation.arguments) try:
result_text = result.to_payload()[:4096] invocation = self._parse_tool_call(payload_text)
except Exception as exc: result = self.tool_registry.execute(invocation.tool_name, invocation.arguments)
result_text = json.dumps({"error": str(exc)[:500]}) result_text = result.to_payload()[:4096]
except Exception as exc:
result_text = json.dumps({"error": str(exc)[:500]})
wrapped = out_start_str + result_text + out_end_str pre_injection_len = len(gen_ids)
# Inject real result tokens into the model's context and the client stream. wrapped = out_start_str + result_text + out_end_str
for rid in self.tokenizer.encode(wrapped): for rid in self.tokenizer.encode(wrapped):
try: try:
rt = self.tokenizer.decode([rid]) rt = self.tokenizer.decode([rid])
except Exception: except Exception:
rt = "" rt = ""
if rt: if rt:
yield "data: " + json.dumps({"token": rt, "gpu": 0}) + "\n\n" yield "data: " + json.dumps({"token": rt, "gpu": 0}) + "\n\n"
_append_token(rid) _append_token(rid)
gen_ids.append(rid) gen_ids.append(rid)
num_generated += 1 num_generated += 1
if num_generated >= max_tokens: if num_generated >= max_tokens:
break break
tool_injected = True tool_injected = True
injection_end_pos = len(gen_ids)
tool_start_idx = -1
# After injection, detect if the model emits ANOTHER full # After injection: if the model starts emitting another <|output_start|>,
# <|output_start|> sequence (training-data loop artifact) and stop the turn. # break the turn — the grounded result already streamed.
if tool_injected and injection_end_pos > 0: elif pre_injection_len > 0 and len(gen_ids) > pre_injection_len + 20:
if _find_subseq(gen_ids, out_start_ids, injection_end_pos) >= 0: try:
post_text = self.tokenizer.decode(gen_ids[pre_injection_len + 10:])
except Exception:
post_text = ""
if out_start_str in post_text:
break break
yield "data: " + json.dumps({"done": True}) + "\n\n" yield "data: " + json.dumps({"done": True}) + "\n\n"