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.562chrm = 'chr1'63CHUNK_SIZE = 1_000_00064fasta_path = '/home/n5huang/dna_token/hg38.fa'65args_token_path = '/home/n5huang/dna_token/output_tokens'66os.makedirs(args_token_path, exist_ok=True)67 68 69 70with pysam.FastaFile(fasta_path) as genome:71    full_sequence = genome.fetch(72            reference=chrm, 73        )74 75print(f"Chromosome: {chrm}")76print(f"Total length: {len(full_sequence):,} bases")77print(f"First 100 bases:\n{full_sequence[:100]}")78 79 80 81# --- 2. LOAD YOUR TOKENIZERS ---82VOCAB_PATHS = {83    "Merged_uni_tfidf":"/home/n5huang/dna_token/tokenizer_evaluation/merge_bpe/merge_tokenizer_unigram_tf_idf.json"84    #"Merged_uni_len":"/home/n5huang/dna_token/tokenizer_evaluation/merge_bpe/merge_tokenizer_unigram_len.json",85    #"Merged_uni_len2":"/home/n5huang/dna_token/tokenizer_evaluation/merge_bpe/merge_tokenizer_unigram_len2.json",86    #"Weighted": "/home/n5huang/dna_token/tokenizer_evaluation/weighted_bpe/tokenizer.json"87    #"SeqOnly": "/home/n5huang/dna_token/tokenizer_evaluation/baseline_bpe/tokenizer.json"88}89 90tokenizers = {}91for name, path in VOCAB_PATHS.items():92    tokenizers[name] = Tokenizer.from_file(path)93 94 95for tok_name, tok in tokenizers.items():96    print(tok.pre_tokenizer)97    print(tok.model)98 99    print(f"\n=== Processing tokenizer: {tok_name} ===")100 101    # 1. Tokenize full chromosome102    raw_tokens = tokenize_full_sequence_collect(tok, full_sequence)103    print(f"Total raw tokens: {len(raw_tokens):,}")104 105    # 2. Build sequences106    final_seqs = nonoverlap_split(107        tokens=raw_tokens,108        maxlen=maxlen,109        tolerance=tolerance110    )111 112    print(f"Total sequences for pretrain: {len(final_seqs):,}")113 114    # 3. Labels115    labels = [chrm] * len(final_seqs)116 117    # 4. Shuffle118    combined = list(zip(final_seqs, labels))119    random.seed(42)120    random.shuffle(combined)121    shuffle_data, shuffle_labels = zip(*combined)122 123    # 5. Train / Val split124    train_num = int(0.9 * len(shuffle_data))125 126    train_data = shuffle_data[:train_num]127    train_labels = shuffle_labels[:train_num]128    val_data = shuffle_data[train_num:]129    val_labels = shuffle_labels[train_num:]130 131    # 6. Save TSVs132    train_path = join(133        args_token_path,134        f"{tok_name}_{chrm}_all_tokenized_train.tsv"135    )136    val_path = join(137        args_token_path,138        f"{tok_name}_{chrm}_all_tokenized_val.tsv"139    )140 141    writetsv(train_data, train_labels, train_path)142    writetsv(val_data, val_labels, val_path)143 144    print(f"Saved:\n  {train_path}\n  {val_path}")145