Spaces:
Running
Running
| # 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() |