| """Explicit biological BPE -> model ID mapping; natural language uses base BPE. |
| |
| Source tokenizers are the historical OmniGene BPEs, with missing single-character |
| alphabet entries appended without altering existing IDs or merges. Every ordinary |
| source piece (including single characters) receives a separate modality-specific |
| model ID. Token display names are serialization metadata, never input text. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
| import shutil |
|
|
| from tokenizers import AddedToken, Tokenizer |
| from transformers import AutoTokenizer |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| DEFAULT = ROOT / 'artifacts/laya_direct_bpe_legacy_representation' |
| LEGACY = Path('/root/autodl-fs/omnigene_v2/scripts/vocab/trained_bpe') |
|
|
|
|
| def digest(path): |
| return hashlib.sha256(Path(path).read_bytes()).hexdigest() |
|
|
|
|
| def prepare(output=DEFAULT, source_root=LEGACY): |
| output = Path(output) |
| if output.exists(): |
| raise FileExistsError(f'Refusing to overwrite representation: {output}') |
| source_root = Path(source_root) |
| source_names = {'dna': 'dna_bpe_20k.json', 'protein': 'protein_bpe_8k.json'} |
| repaired_sources, repairs = {}, {} |
| for kind,name in source_names.items(): |
| original = json.loads((source_root/name).read_text()) |
| vocab = original['model']['vocab'] |
| |
| |
| alphabet = 'ACGTN' if kind=='dna' else 'ABCDEFGHIJKLMNOPQRSTUVWXYZ' |
| missing = [letter for letter in alphabet if letter not in vocab] |
| for letter in missing: |
| vocab[letter] = max(vocab.values())+1 |
| repairs[kind] = missing |
| repaired_sources[kind] = Tokenizer.from_str(json.dumps(original)) |
| base = AutoTokenizer.from_pretrained(ROOT / 'artifacts/laya_model/tokenizer') |
| expanded = AutoTokenizer.from_pretrained(ROOT / 'artifacts/laya_model/tokenizer') |
| entries = [] |
| for kind, name in source_names.items(): |
| source = repaired_sources[kind] |
| for piece, source_id in sorted(source.get_vocab().items(), key=lambda pair: pair[1]): |
| if piece in {'[UNK]', '[PAD]'}: |
| continue |
| if not piece.isalpha(): |
| raise ValueError(f'Unexpected source vocabulary piece: {piece!r}') |
| ids = base(piece, add_special_tokens=False)['input_ids'] |
| if not ids or any(i in base.all_special_ids for i in ids): |
| raise ValueError(f'Cannot initialize {kind}:{piece} from raw fragment') |
| entries.append({'modality': kind, 'source_id': source_id, 'piece': piece, |
| 'token': f'<BIO_{kind.upper()}:{piece}>', 'base_ids': ids}) |
| added = expanded.add_tokens([AddedToken(x['token'], normalized=False) for x in entries]) |
| if added != len(entries): |
| raise ValueError('Duplicate or conflicting vocabulary entries') |
| vocab = expanded.get_vocab() |
| for entry in entries: |
| entry['expanded_id'] = vocab[entry['token']] |
| if any(vocab[k] != v for k, v in base.get_vocab().items()): |
| raise ValueError('Original token IDs changed') |
| output.mkdir(parents=True) |
| (output/'source_tokenizers').mkdir() |
| (output/'original_source_tokenizers').mkdir() |
| for kind,name in source_names.items(): |
| repaired_sources[kind].save(str(output/'source_tokenizers'/name)) |
| shutil.copy2(source_root/name, output/'original_source_tokenizers'/name) |
| scripts = source_root.parent |
| provenance_scripts = [scripts/name for name in ('1-sample_bpe_corpus.py','2-train_bio_bpe.py','3-expand_tokenizer.py')] |
| (output/'provenance_scripts').mkdir() |
| for path in provenance_scripts: |
| if path.exists():shutil.copy2(path,output/'provenance_scripts'/path.name) |
| base.save_pretrained(output / 'base_tokenizer') |
| expanded.save_pretrained(output / 'expanded_tokenizer') |
| (output / 'new_tokens.json').write_text(json.dumps(entries, ensure_ascii=False, indent=2)+'\n') |
| meta = {'schema_version': 'laya-direct-bpe-legacy-v1', 'no_cpt': True, |
| 'fit_policy': 'reuse external historical BPE; no refitting; no neural CPT', |
| 'historical_fit_recipe': 'sample about 1 GiB each from dna_32g.txt and protein_uni_16.txt, seed42; BPE min_frequency10', |
| 'historical_fit_corpus_overlap_with_evaluation': 'not established; no train-only claim', |
| 'source_files': {k:{'path':str(source_root/n),'sha256':digest(source_root/n)} for k,n in source_names.items()}, |
| 'alphabet_repair': repairs, 'existing_merges_preserved':True, |
| 'old_vocab_size': len(base), 'new_vocab_size': len(expanded), |
| 'added_by_modality': {k: sum(x['modality'] == k for x in entries) for k in source_names}, |
| 'initialization': 'mean of original tokenizer IDs of RAW biological fragment; no delimiters', |
| 'encoding': 'source BPE ID mapped directly to added model ID; natural language uses base tokenizer', |
| 'output_hashes': {str(p.relative_to(output)): digest(p) for p in sorted(output.rglob('*')) if p.is_file()}} |
| (output / 'metadata.json').write_text(json.dumps(meta, indent=2)+'\n') |
| print(json.dumps({k: meta[k] for k in ('old_vocab_size','new_vocab_size','added_by_modality')}, indent=2)) |
|
|
|
|
| class DirectBPE: |
| def __init__(self, directory=DEFAULT): |
| self.root = Path(directory) |
| self.meta = json.loads((self.root / 'metadata.json').read_text()) |
| for name, expected in self.meta['output_hashes'].items(): |
| if digest(self.root / name) != expected: |
| raise ValueError(f'Representation hash mismatch: {name}') |
| self.base = AutoTokenizer.from_pretrained(self.root / 'base_tokenizer') |
| self.expanded = AutoTokenizer.from_pretrained(self.root / 'expanded_tokenizer') |
| self.entries = json.loads((self.root / 'new_tokens.json').read_text()) |
| self.sources = {k: Tokenizer.from_file(str(self.root / 'source_tokenizers' / name)) |
| for k, name in [('dna','dna_bpe_20k.json'),('protein','protein_bpe_8k.json')]} |
| self.maps = {k: {e['source_id']: e['expanded_id'] for e in self.entries if e['modality'] == k} |
| for k in self.sources} |
| self.inverse = {e['expanded_id']: e['piece'] for e in self.entries} |
|
|
| def sequence_ids(self, sequence, kind): |
| encoded = self.sources[kind].encode(sequence) |
| if ''.join(encoded.tokens) != sequence: |
| raise ValueError('Source BPE changed the sequence or produced UNK') |
| try: |
| ids = [self.maps[kind][i] for i in encoded.ids] |
| except KeyError as exc: |
| raise ValueError('Unmapped source BPE token') from exc |
| if ''.join(self.inverse[i] for i in ids) != sequence: |
| raise ValueError('Model IDs do not reconstruct the original sequence') |
| return ids |
|
|
| def build(self, row, condition, builder, order, max_len=1024, head_max_len=256): |
| q = {'t': 'choice', 'ins': 'Choose the correct biological label for the sequence.', |
| 'crit': {x: None for x in row['choices']}} |
| prefix = f"Task context: {row['context']}\nSequence: " |
| if condition == 'raw': |
| ids, markers = builder(self.base, prefix + row['sequence'], q, |
| max_len=1_000_000, head_max_len=head_max_len, option_order=order) |
| elif condition == 'full_bpe': |
| |
| |
| |
| header, markers = builder(self.base, '', q, max_len=1_000_000, |
| head_max_len=head_max_len, option_order=order) |
| context_ids = self.base(prefix.replace(self.base.mask_token, ' '), add_special_tokens=False)['input_ids'] |
| ids = header[:-1] + context_ids + self.sequence_ids(row['sequence'], row['kind']) + [self.base.sep_token_id] |
| else: |
| raise ValueError(condition) |
| if len(ids) > max_len: |
| raise ValueError(f"Full input exceeds {max_len}: {row['id']} {condition} {len(ids)}") |
| if len(markers) != len(row['choices']): |
| raise ValueError('Candidate marker loss') |
| for canonical, marker in zip(order, markers): |
| expected = self.base(' '+row['choices'][canonical], add_special_tokens=False)['input_ids'] |
| if ids[marker+1:marker+1+len(expected)] != expected: |
| raise ValueError('Candidate text truncated or remapped incorrectly') |
| return ids, markers |
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('--output', type=Path, default=DEFAULT) |
| parser.add_argument('--source-dir', type=Path, default=LEGACY) |
| args=parser.parse_args() |
| prepare(args.output,args.source_dir) |
|
|
