Add MP-01 (9500) 100M base: weights, config, tokenizer, model code, model card and eval results
11e2c7c verified | """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} | |
| 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):] | |