# ---------- 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")