File size: 3,846 Bytes
655f931 | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | #!/usr/bin/env python3
"""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) # remove stage directions
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()
|