Spaces:
Running
Running
File size: 4,699 Bytes
a4c4002 | 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 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | # Updated for ShareGPT
import os
import pandas as pd
import numpy as np
from transformers import AutoTokenizer
import matplotlib.pyplot as plt # Optional, for distribution logic
# ================= CONFIGURATION =================
# Path to your unified parquet file (created by PMPF3)
DATASET_FILE = os.path.join(os.getcwd(), "dataset_cache", "unified_dataset.parquet")
# Path to your model (for the tokenizer)
MODEL_PATH = os.getcwd()
# =================================================
# =================================================
# Mistral Tekken Templates for exact length scanning
MT_SYSTEM = "[SYSTEM_PROMPT]{sys}[/SYSTEM_PROMPT]"
MT_USER = "[INST] {q} [/INST]"
MT_ASSISTANT = "{a}</s>"
def format_sharegpt_to_text(conversations):
"""
Mimics how the model sees the data during training.
Reconstructs the full sequence text from the structured ShareGPT list.
"""
if not isinstance(conversations, (list, np.ndarray)):
return ""
text = "<s>"
for turn in conversations:
if not isinstance(turn, dict):
continue
role = turn.get('from', '')
content = turn.get('value', '')
if role == 'system':
text += MT_SYSTEM.format(sys=content)
elif role == 'human':
text += MT_USER.format(q=content)
elif role == 'gpt':
text += MT_ASSISTANT.format(a=content)
return text
def main():
print("--- DATASET CONTEXT SCANNER ---")
# 1. Load Tokenizer
print(f"Loading Tokenizer from {MODEL_PATH}...")
try:
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, local_files_only=True)
except Exception as e:
print(f"Error loading tokenizer: {e}")
print("Cannot calculate exact tokens without tokenizer.")
return
# 2. Load Dataset
if not os.path.exists(DATASET_FILE):
print(f"Error: {DATASET_FILE} not found.")
print("Run the Unify script (PMPF3) first.")
return
print(f"Loading dataset: {DATASET_FILE}")
df = pd.read_parquet(DATASET_FILE)
print(f"Scanning {len(df)} entries...")
# 3. Scan
lengths = []
char_counts = []
max_tokens = 0
max_chars = 0
longest_text = ""
# Iterate and measure
for index, row in df.iterrows():
# Retrieve the structured conversations array
if 'conversations' not in row or row['conversations'] is None:
continue
conversations = row['conversations']
# Reconstruct the string sequence as tokenized on GPU
text = format_sharegpt_to_text(conversations)
if not text:
continue
# Count characters
c_len = len(text)
char_counts.append(c_len)
# Count tokens (Exact)
# We use len(encode) to get the real number the GPU sees
ids = tokenizer.encode(text, add_special_tokens=True)
t_len = len(ids)
lengths.append(t_len)
if t_len > max_tokens:
max_tokens = t_len
max_chars = c_len
longest_text = text
# 4. Statistics
lengths = np.array(lengths)
p95 = np.percentile(lengths, 95)
p99 = np.percentile(lengths, 99)
print("\n" + "="*40)
print(" SCAN RESULTS ")
print("="*40)
print(f"Total Entries: {len(df)}")
print(f"Max Characters: {max_chars}")
print(f"Max Tokens (Exact): {max_tokens}")
print("-" * 40)
print(f"Average Tokens: {int(np.mean(lengths))}")
print(f"95th Percentile: {int(p95)} tokens (Covers 95% of data)")
print(f"99th Percentile: {int(p99)} tokens (Covers 99% of data)")
print("="*40)
# 5. Recommendation
print("\n--- RECOMMENDATION ---")
# Round up to nearest multiple of 64 for GPU efficiency
recommended_ctx = int(max_tokens)
remainder = recommended_ctx % 64
if remainder != 0:
recommended_ctx += (64 - remainder)
# Add a tiny safety buffer
recommended_ctx += 64
print(f"Set MAX_SEQ_LENGTH = {recommended_ctx}")
if recommended_ctx > 2048:
print("\n⚠️ WARNING: This is high for an 8GB card.")
print(f"Consider truncating to {int(p99)} (99th percentile) to save VRAM.")
else:
print(f"✅ Safe for 8GB VRAM (assuming LoRA r=16).")
print("\n--- LONGEST ENTRY PREVIEW ---")
print(longest_text[:500] + "..." + longest_text[-100:])
if __name__ == "__main__":
main() |