Md Najmul Hossain
Upload python_scripts/build_token_cache_mkd_auto_resume_final.py with huggingface_hub
3204bda verified | # ---------- build_token_cache_mkd_auto_resume_v6.py ---------- | |
| # (์ฃผ์ํ์ฌ์ ์ผ์ด๋ | MKD) | |
| # Fully automated, resumable, logging-enabled token cache system | |
| # Includes daily logs, auto PKL/TMP rotation, and safe checkpoints | |
| import os | |
| import time | |
| import pickle | |
| import psutil | |
| import subprocess | |
| import sys | |
| import torch | |
| import gc | |
| from datetime import datetime | |
| from tokenizers import Tokenizer | |
| # ============================================================ | |
| # CONFIGURATION | |
| # ============================================================ | |
| BASE_DIR = "/home/mkd-s003/AI/finetune" | |
| TOKENIZER_PATH = os.path.join(BASE_DIR, "prod_tokenizer_v2.json") | |
| TEXT_PATH = os.path.join(BASE_DIR, "clean_text.txt") | |
| TRAIN_PKL = os.path.join(BASE_DIR, "token_cache_train.pkl") | |
| VAL_PKL = os.path.join(BASE_DIR, "token_cache_val.pkl") | |
| TMP_PKL = os.path.join(BASE_DIR, "token_cache_progress.tmp") | |
| LOG_FILE = os.path.join(BASE_DIR, "token_cache_log.txt") | |
| TOKEN_LIMIT_TRAIN = 200_000_000 | |
| TOKEN_LIMIT_VAL = 20_000_000 | |
| SPLIT_RATIO = 0.9 | |
| SAVE_EVERY = 500_000 | |
| BATCH_LINES = 3_000_000 | |
| RESTART_DELAY = 60 | |
| # ============================================================ | |
| # LOGGING UTILITY | |
| # ============================================================ | |
| def log_message(msg): | |
| """Write timestamped message to both console and log file""" | |
| stamp = datetime.now().strftime("[%Y-%m-%d %H:%M:%S]") | |
| text = f"{stamp} {msg}" | |
| print(text) | |
| with open(LOG_FILE, "a", encoding="utf-8") as f: | |
| f.write(text + "\n") | |
| # ============================================================ | |
| # MONITOR UTILITIES | |
| # ============================================================ | |
| def system_status(): | |
| cpu = psutil.cpu_percent(interval=0.2) | |
| mem = psutil.virtual_memory().percent | |
| gpu = "" | |
| if torch.cuda.is_available(): | |
| try: | |
| util = subprocess.check_output( | |
| ["nvidia-smi", "--query-gpu=utilization.gpu", | |
| "--format=csv,noheader,nounits"] | |
| ).decode().strip().split("\n")[0] | |
| gpu = f" | GPU Util {util}%" | |
| except Exception: | |
| gpu = "" | |
| return f"CPU {cpu:.0f}% | RAM {mem:.1f}%{gpu}" | |
| # ============================================================ | |
| # BACKUP MANAGEMENT | |
| # ============================================================ | |
| def rotate_backups(): | |
| timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") | |
| for f in [TRAIN_PKL, VAL_PKL, TMP_PKL]: | |
| if os.path.exists(f): | |
| new_name = f.replace(".pkl", f"_{timestamp}.bak").replace(".tmp", f"_{timestamp}.bak") | |
| os.rename(f, new_name) | |
| log_message(f"Old file backed up โ {new_name}") | |
| # Keep only latest backup | |
| for file in sorted(os.listdir(BASE_DIR)): | |
| if file.endswith(".bak"): | |
| file_path = os.path.join(BASE_DIR, file) | |
| age_hours = (time.time() - os.path.getmtime(file_path)) / 3600 | |
| if age_hours > 24: # older than 1 day | |
| os.remove(file_path) | |
| log_message(f"Deleted old backup โ {file}") | |
| # ============================================================ | |
| # MAIN FUNCTION | |
| # ============================================================ | |
| def build_token_cache(): | |
| log_message("===== Starting MKD Token Cache Builder v6 =====") | |
| log_message(f"Tokenizer: {TOKENIZER_PATH}") | |
| log_message(f"Dataset: {TEXT_PATH}") | |
| tokenizer = Tokenizer.from_file(TOKENIZER_PATH) | |
| eot = tokenizer.token_to_id("<|endoftext|>") | |
| log_message(f"Tokenizer loaded successfully. EOT id = {eot}") | |
| # Resume progress | |
| start_line = 0 | |
| if os.path.exists(TMP_PKL): | |
| with open(TMP_PKL, "rb") as f: | |
| saved = pickle.load(f) | |
| start_line = saved.get("line", 0) | |
| log_message(f"Resuming from line {start_line:,}") | |
| else: | |
| log_message("No checkpoint found. Starting fresh.") | |
| rotate_backups() | |
| # Count dataset | |
| with open(TEXT_PATH, "r", encoding="utf-8") as f: | |
| total_lines = sum(1 for _ in f) | |
| split_index = int(total_lines * SPLIT_RATIO) | |
| log_message(f"Total lines: {total_lines:,} | Train split @ {split_index:,}") | |
| train_file = open(TRAIN_PKL, "ab") | |
| val_file = open(VAL_PKL, "ab") | |
| start_time = time.time() | |
| total_train_tokens, total_val_tokens = 0, 0 | |
| batch_lines = [] | |
| with open(TEXT_PATH, "r", encoding="utf-8") as f: | |
| for i, line in enumerate(f): | |
| if i < start_line: | |
| continue | |
| line = line.strip() | |
| if not line: | |
| continue | |
| batch_lines.append(line) | |
| if len(batch_lines) >= BATCH_LINES: | |
| t_tok, v_tok = process_batch(tokenizer, batch_lines, eot, i, total_lines, | |
| split_index, train_file, val_file, start_time) | |
| total_train_tokens += t_tok | |
| total_val_tokens += v_tok | |
| batch_lines = [] | |
| gc.collect() | |
| # token limit reached | |
| if total_train_tokens >= TOKEN_LIMIT_TRAIN and total_val_tokens >= TOKEN_LIMIT_VAL: | |
| log_message("Target token limit reached โ stopping safely.") | |
| break | |
| if (i + 1) % SAVE_EVERY == 0: | |
| with open(TMP_PKL, "wb") as tmp: | |
| pickle.dump({"line": i}, tmp) | |
| log_message(f"Checkpoint saved at line {i:,} | {system_status()}") | |
| if batch_lines: | |
| t_tok, v_tok = process_batch(tokenizer, batch_lines, eot, total_lines, total_lines, | |
| split_index, train_file, val_file, start_time) | |
| total_train_tokens += t_tok | |
| total_val_tokens += v_tok | |
| train_file.close() | |
| val_file.close() | |
| # Cleanup + new tmp | |
| if os.path.exists(TMP_PKL): | |
| os.remove(TMP_PKL) | |
| with open(TMP_PKL, "wb") as tmp: | |
| pickle.dump({"line": i}, tmp) | |
| log_message(f"New progress checkpoint saved at line {total_lines:,}") | |
| # Summary | |
| elapsed_hours = (time.time() - start_time) / 3600 | |
| log_message("========== SUMMARY ==========") | |
| log_message(f"Train tokens generated: {total_train_tokens:,}") | |
| log_message(f"Val tokens generated: {total_val_tokens:,}") | |
| log_message(f"Train cache โ {TRAIN_PKL}") | |
| log_message(f"Val cache โ {VAL_PKL}") | |
| log_message(f"Elapsed time: {elapsed_hours:.2f} hours") | |
| log_message("=============================") | |
| # ============================================================ | |
| # PROCESS BATCH | |
| # ============================================================ | |
| def process_batch(tokenizer, lines, eot, current_line, total_lines, | |
| split_index, train_file, val_file, start_time): | |
| encoded = tokenizer.encode_batch(lines) | |
| train_tokens, val_tokens = [], [] | |
| for j, enc in enumerate(encoded): | |
| line_idx = current_line - len(encoded) + j + 1 | |
| if line_idx < split_index: | |
| train_tokens.append(enc.ids + [eot]) | |
| else: | |
| val_tokens.append(enc.ids + [eot]) | |
| train_tok_count = sum(len(x) for x in train_tokens) | |
| val_tok_count = sum(len(x) for x in val_tokens) | |
| for seq in train_tokens: | |
| pickle.dump(seq, train_file) | |
| for seq in val_tokens: | |
| pickle.dump(seq, val_file) | |
| elapsed = (time.time() - start_time) / 60 | |
| percent = (current_line / total_lines) * 100 | |
| remaining = (elapsed / (percent + 1e-5)) * (100 - percent) | |
| msg = (f"Progress {percent:5.2f}% | Line {current_line:,}/{total_lines:,} " | |
| f"| TrainTok +{train_tok_count:,} | ValTok +{val_tok_count:,} " | |
| f"| Elapsed {elapsed:.1f}m | ETA {remaining/60:.1f}h | {system_status()}") | |
| log_message(msg) | |
| return train_tok_count, val_tok_count | |
| # ============================================================ | |
| # AUTO RESTART LOOP | |
| # ============================================================ | |
| if __name__ == "__main__": | |
| while True: | |
| try: | |
| build_token_cache() | |
| break | |
| except KeyboardInterrupt: | |
| log_message("Stopped manually. Progress saved.") | |
| sys.exit(0) | |
| except Exception as e: | |
| log_message(f"Error: {e}") | |
| log_message(f"Restarting automatically in {RESTART_DELAY}s...") | |
| for t in range(RESTART_DELAY, 0, -1): | |
| print(f"\rRestart in {t:02d}s | {system_status()}", end="", flush=True) | |
| time.sleep(1) | |
| print("\nRestarting now...\n") | |