File size: 3,884 Bytes
66dee2d | 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 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | #!/usr/bin/env python3
"""
Analyze token-level sequence lengths in aes_all.jsonl using the actual tokenizer.
Applies the chat template to get real token counts.
Outputs statistics and a recommended max_seq_length.
"""
import json
import sys
from pathlib import Path
from transformers import AutoTokenizer
WORKSPACE = Path("/workspace/elinnos")
INPUT = WORKSPACE / "aes_all.jsonl"
TOKENIZER_SRC = WORKSPACE / "Qwen2.5-7B-Instruct"
CHAT_TEMPLATE = WORKSPACE / "elinnos-qwen2.5-7b-multi-ip-lora-v4" / "chat_template.jinja"
def main():
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(str(TOKENIZER_SRC), trust_remote_code=True)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
if CHAT_TEMPLATE.is_file():
tokenizer.chat_template = CHAT_TEMPLATE.read_text()
print(f"Chat template loaded from {CHAT_TEMPLATE}")
print(f"Loading {INPUT}...")
records = []
with INPUT.open("r", encoding="utf-8") as f:
for line in f:
line = line.strip()
if line:
records.append(json.loads(line))
print(f"Total samples: {len(records)}")
print("Tokenizing with chat template (this may take a minute)...")
token_lengths = []
for i, rec in enumerate(records):
messages = rec["messages"]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
tokens = tokenizer(text, return_tensors=None, add_special_tokens=False)
token_lengths.append(len(tokens["input_ids"]))
if (i + 1) % 500 == 0:
print(f" Processed {i+1}/{len(records)}...")
token_lengths.sort()
n = len(token_lengths)
print(f"\n{'='*60}")
print(" TOKEN LENGTH STATISTICS")
print(f"{'='*60}")
print(f" Samples: {n}")
print(f" Min: {token_lengths[0]}")
print(f" Max: {token_lengths[-1]}")
print(f" Mean: {sum(token_lengths)/n:.0f}")
print(f" Median: {token_lengths[n//2]}")
for p in [75, 90, 95, 97, 99, 100]:
idx = min(int(n * p / 100), n - 1)
print(f" {p}th pct: {token_lengths[idx]}")
# Recommend seq length
p95 = token_lengths[min(int(n * 0.95), n - 1)]
p99 = token_lengths[min(int(n * 0.99), n - 1)]
p100 = token_lengths[-1]
if p95 <= 4096:
recommended = 4096
elif p95 <= 8192:
recommended = 8192
elif p99 <= 12288:
recommended = 12288
else:
recommended = min(p100, 16384)
# Round up to nearest 512
recommended = ((recommended + 511) // 512) * 512
print(f"\n Recommended max_seq_length: {recommended}")
print(f" (covers 95th pct={p95}, 99th pct={p99}, max={p100})")
# Count how many samples would be truncated at various lengths
for sl in [4096, 8192, 12288, 16384]:
truncated = sum(1 for t in token_lengths if t > sl)
print(f" At seq_len={sl}: {truncated}/{n} ({truncated/n*100:.1f}%) samples truncated")
out_path = WORKSPACE / "aes_training" / "data" / "recommended_seq_length.txt"
out_path.parent.mkdir(parents=True, exist_ok=True)
out_path.write_text(str(recommended))
print(f"\n Written to: {out_path}")
# Also save full stats
stats_path = WORKSPACE / "aes_training" / "data" / "seq_length_stats.json"
stats = {
"total_samples": n,
"min": token_lengths[0],
"max": token_lengths[-1],
"mean": sum(token_lengths) / n,
"median": token_lengths[n // 2],
"percentiles": {
str(p): token_lengths[min(int(n * p / 100), n - 1)]
for p in [50, 75, 90, 95, 97, 99, 100]
},
"recommended_max_seq_length": recommended,
}
with stats_path.open("w") as f:
json.dump(stats, f, indent=2)
print(f" Stats saved to: {stats_path}")
if __name__ == "__main__":
main()
|