File size: 3,038 Bytes
11e2c7c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | """Shared ChatML tokens, assistant supervision and context budgeting."""
import os
from functools import lru_cache
from pathlib import Path
from tokenizers import Tokenizer
ROOT = Path(__file__).resolve().parent.parent
SPECIALS = ('<|endoftext|>', '<|pad|>', '<|im_start|>', '<|im_end|>', '<|system|>', '<|user|>', '<|assistant|>', '<|reserved|>')
STOP_IDS = {0, 3}
@lru_cache(maxsize=1)
def tokenizer():
return Tokenizer.from_file(str(ROOT / os.environ.get('VBQ_TOKENIZER', 'data/tokenizer/vbq-tokenizer.json')))
def render(messages, add_generation_prompt=False):
ids, mask = [], []
tok = tokenizer()
for m in messages:
role, content = m['role'], m['content']
if role not in ('system', 'user', 'assistant') or not isinstance(content, str):
raise ValueError('messages must contain system/user/assistant roles and string content')
if any(s in content for s in SPECIALS):
raise ValueError('message content cannot contain reserved tokens')
header = tok.encode(f'<|im_start|>{role}\n', add_special_tokens=False).ids
body = tok.encode(content, add_special_tokens=False).ids
tail = tok.encode('\n', add_special_tokens=False).ids
ids.extend(header + body + [3] + tail)
mask.extend([0] * len(header) + [int(role == 'assistant')] * (len(body) + 1) + [0] * len(tail))
if add_generation_prompt:
header = tok.encode('<|im_start|>assistant\n', add_special_tokens=False).ids
ids.extend(header); mask.extend([0] * len(header))
return ids, mask
def fit_history(messages, budget):
history = [dict(m) for m in messages]
if not history or history[-1]['role'] != 'user':
raise ValueError('chat must end with a user message')
if any(m['role'] == 'system' for m in history[1:]):
raise ValueError('system message must be first')
start = int(history[0]['role'] == 'system')
while True:
ids, _ = render(history, add_generation_prompt=True)
if len(ids) <= budget: return history, ids
if len(history) <= start + 1:
raise ValueError('system prompt and latest message exceed the context budget')
end = start + 1
while end < len(history) - 1 and history[end]['role'] != 'user': end += 1
del history[start:end]
def stream_reply(model, messages, max_new=256, temperature=0.7, top_p=0.9, top_k=50, rep_penalty=1.1):
import torch
_, ids = fit_history(messages, model.c.max_seq_len - max_new)
out, sent = [], ''
for t in model.generate_cached(torch.tensor([ids], device=next(model.parameters()).device), max_new_tokens=max_new,
temperature=temperature, top_p=top_p, top_k=top_k, repetition_penalty=rep_penalty, eos_id=STOP_IDS):
if t in STOP_IDS: break
out.append(t); text = tokenizer().decode(out)
if text.endswith('�'): continue
yield text[len(sent):]; sent = text
text = tokenizer().decode(out)
if len(text) > len(sent): yield text[len(sent):]
|