nancyH/token_evaluation

1import os2import random3from os.path import join4from collections import Counter5import numpy as np6import pysam7from tokenizers import Tokenizer8from tokenizers.models import BPE9from tokenizers.trainers import BpeTrainer10from tokenizers.pre_tokenizers import PreTokenizer11from tokenizers.pre_tokenizers import ByteLevel12from tokenizers.pre_tokenizers import Whitespace13from tokenizers.pre_tokenizers import CharDelimiterSplit14from tokenizers.normalizers import Sequence, Lowercase15from tokenizers import models, pre_tokenizers, decoders16from tokenizers.pre_tokenizers import Split17 18 19def writetsv(data, label, savefile):20    with open(savefile, 'w') as f:21        f.write('sequence\tlabels\n')22        for seq, lab in zip(data, label):23            f.write(f'{seq}\t{lab}\n')24 25 26def nonoverlap_split(tokens, maxlen, tolerance=0.5):27    seqs = []28    skipped = 029 30    num_windows = len(tokens) // maxlen31 32    for i in range(num_windows):33        window = tokens[i*maxlen:(i+1)*maxlen]34 35        # NEW: token-aware N detection36        num_N = sum('N' in tok for tok in window)37 38        if num_N / maxlen < tolerance:39            seqs.append(" ".join(window))40        else:41            skipped += 142 43    print(f"In this chromosome, skipped sequences: {skipped}")44    return seqs45 46 47def tokenize_full_sequence_collect(tokenizer, sequence, chunk_size=1_000_000):48    raw_tokens = []49 50    for i in range(0, len(sequence), chunk_size):51        chunk = sequence[i:i + chunk_size]52        encoded = tokenizer.encode(chunk)53        raw_tokens.extend(encoded.tokens)54 55        if i % (10 * chunk_size) == 0:56            print(f"Processed {i:,} bp")57 58    return raw_tokens59 60maxlen = 51261tolerance = 0.562CHUNK_SIZE = 1_000_00063fasta_path = '/home/n5huang/dna_token/hg38.fa'64args_token_path = '/home/n5huang/dna_token/output_tokens'65os.makedirs(args_token_path, exist_ok=True)66 67# Set to a list like ["chr1", "chr2"] to limit processing.68CHROMOSOMES = None69EXCLUDE_CHROMS = set()70 71 72# --- ABLATION TOKENIZERS ---73VOCAB_PATHS = {74    "ablation_no_partition": "/home/n5huang/dna_token/tokenizer_evaluation/ablation/vocab_5120/ablation_no_partition.json",75    "ablation_no_priority":  "/home/n5huang/dna_token/tokenizer_evaluation/ablation/vocab_5120/ablation_no_priority.json",76    "ablation_no_length":    "/home/n5huang/dna_token/tokenizer_evaluation/ablation/vocab_5120/ablation_no_length.json",77}78 79tokenizers = {}80for name, path in VOCAB_PATHS.items():81    tokenizers[name] = Tokenizer.from_file(path)82 83 84with pysam.FastaFile(fasta_path) as genome:85    chroms = genome.references if CHROMOSOMES is None else CHROMOSOMES86    chroms = [c for c in chroms if c not in EXCLUDE_CHROMS]87 88    for tok_name, tok in tokenizers.items():89        print(tok.pre_tokenizer)90        print(tok.model)91 92        print(f"\n=== Processing tokenizer: {tok_name} ===")93 94        all_seqs = []95        all_labels = []96 97        for chrm in chroms:98            full_sequence = genome.fetch(reference=chrm)99 100            print(f"\nChromosome: {chrm}")101            print(f"Total length: {len(full_sequence):,} bases")102            print(f"First 100 bases:\n{full_sequence[:100]}")103 104            # 1. Tokenize full chromosome105            raw_tokens = tokenize_full_sequence_collect(106                tok,107                full_sequence,108                chunk_size=CHUNK_SIZE109            )110            print(f"Total raw tokens: {len(raw_tokens):,}")111 112            # 2. Build sequences113            final_seqs = nonoverlap_split(114                tokens=raw_tokens,115                maxlen=maxlen,116                tolerance=tolerance117            )118 119            print(f"Total sequences for pretrain: {len(final_seqs):,}")120 121            all_seqs.extend(final_seqs)122            all_labels.extend([chrm] * len(final_seqs))123 124        if not all_seqs:125            print(f"No sequences generated for tokenizer: {tok_name}")126            continue127 128        # 3. Shuffle129        combined = list(zip(all_seqs, all_labels))130        random.seed(42)131        random.shuffle(combined)132        shuffle_data, shuffle_labels = zip(*combined)133 134        # 4. Train / Val split135        train_num = int(0.9 * len(shuffle_data))136 137        train_data = shuffle_data[:train_num]138        train_labels = shuffle_labels[:train_num]139        val_data = shuffle_data[train_num:]140        val_labels = shuffle_labels[train_num:]141 142        # 5. Save TSVs143        train_path = join(144            args_token_path,145            f"{tok_name}_allchr_all_tokenized_train.tsv"146        )147        val_path = join(148            args_token_path,149            f"{tok_name}_allchr_all_tokenized_val.tsv"150        )151 152        writetsv(train_data, train_labels, train_path)153        writetsv(val_data, val_labels, val_path)154 155        print(f"Saved:\n  {train_path}\n  {val_path}")156