| |
| """Prepare TinyDialogues for a small chat-style BitNet run.""" |
| import argparse, json, re |
| from pathlib import Path |
| import sentencepiece as spm |
|
|
| ROLE_RE = re.compile(r"\*\*(.+?)\*\*:\s*(.*?)(?=\*\*|$)", re.S) |
|
|
| def clean(s: str) -> str: |
| s = s.replace('\\n', '\n') |
| s = re.sub(r'\*[^*]*\*', ' ', s) |
| s = s.strip().strip('"').strip("'") |
| s = re.sub(r'\s+', ' ', s).strip() |
| return s |
|
|
| def parse(line: str): |
| line = line.replace('\\n', '\n') |
| turns = [] |
| for role, text in ROLE_RE.findall(line): |
| role = role.strip().lower() |
| text = clean(text) |
| if not text: |
| continue |
| speaker = 'user' if 'child' in role else 'assistant' |
| if turns and turns[-1][0] == speaker: |
| turns[-1][1] += ' ' + text |
| else: |
| turns.append([speaker, text]) |
| return turns |
|
|
| def fmt(turns): |
| return ' '.join(f"{'User' if r=='user' else 'Assistant'}: {t}" for r, t in turns).strip() |
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--train', required=True) |
| ap.add_argument('--val', required=True) |
| ap.add_argument('--out', required=True) |
| ap.add_argument('--vocab-size', type=int, default=2048) |
| args = ap.parse_args() |
| out = Path(args.out); out.mkdir(parents=True, exist_ok=True) |
| stats = {} |
| for split, src in [('train', args.train), ('val', args.val)]: |
| dst = out / f'formatted_{split}.txt' |
| n = skipped = 0 |
| with open(src, encoding='utf-8') as f, open(dst, 'w', encoding='utf-8') as g: |
| for line in f: |
| turns = parse(line) |
| if len(turns) < 2 or not any(r == 'assistant' for r, _ in turns): |
| skipped += 1 |
| continue |
| g.write(fmt(turns) + '\n') |
| n += 1 |
| stats[split] = {'dialogues': n, 'skipped': skipped, 'file': str(dst)} |
| print(split, stats[split], flush=True) |
|
|
| prefix = out / 'tokenizer' |
| spm.SentencePieceTrainer.train( |
| input=str(out / 'formatted_train.txt'), |
| model_prefix=str(prefix), |
| vocab_size=args.vocab_size, |
| model_type='bpe', |
| character_coverage=1.0, |
| byte_fallback=True, |
| unk_id=0, bos_id=1, eos_id=2, pad_id=3, |
| unk_piece='<unk>', bos_piece='<s>', eos_piece='</s>', pad_piece='<pad>', |
| input_sentence_size=2000000, |
| shuffle_input_sentence=True, |
| max_sentence_length=100000, |
| normalization_rule_name='nmt_nfkc_cf', |
| ) |
| sp = spm.SentencePieceProcessor(model_file=str(prefix) + '.model') |
| stats['vocab_size'] = sp.get_piece_size() |
| chat_template = ( |
| "{% for message in messages %}" |
| "{% if message['role'] == 'user' %}User: {{ message['content'] }} " |
| "{% elif message['role'] == 'assistant' %}Assistant: {{ message['content'] }} " |
| "{% endif %}" |
| "{% endfor %}" |
| "{% if add_generation_prompt %}Assistant:{% endif %}" |
| ) |
| tok_cfg = { |
| 'tokenizer_class': 'LlamaTokenizer', |
| 'model_max_length': 128, |
| 'bos_token': '<s>', 'eos_token': '</s>', 'pad_token': '<pad>', 'unk_token': '<unk>', |
| 'add_bos_token': True, 'add_eos_token': False, |
| 'clean_up_tokenization_spaces': True, |
| 'chat_template': chat_template, |
| } |
| (out / 'tokenizer_config.json').write_text(json.dumps(tok_cfg, indent=2), encoding='utf-8') |
| (out / 'special_tokens_map.json').write_text(json.dumps({ |
| 'bos_token': '<s>', 'eos_token': '</s>', 'pad_token': '<pad>', 'unk_token': '<unk>' |
| }, indent=2), encoding='utf-8') |
| (out / 'chat_template.jinja').write_text(chat_template, encoding='utf-8') |
| (out / 'stats.json').write_text(json.dumps(stats, indent=2), encoding='utf-8') |
| print(json.dumps(stats, indent=2)) |
|
|
| if __name__ == '__main__': |
| main() |
|
|
