#!/usr/bin/env python3 """ Domain-Adaptive MLM Pretrainer Continues pretraining BERT-family models (BERT, RoBERTa, DeBERTa, ELECTRA, RemBERT, etc.) on an unlabelled domain corpus using Masked Language Modelling (MLM). Outputs a saved encoder checkpoint per model that can be referenced directly as `model_name` in ensemble_config.json for downstream fine-tuning. Usage: python mlm_pretrainer.py \ --data_path listings.parquet \ --text_col text \ --pretrain_config_path pretrain_config.json \ --output_base_dir pretrained_checkpoints """ import signal import atexit import argparse import gc import json import warnings import sys import copy import shutil import re from datetime import datetime from pathlib import Path import numpy as np import torch from datasets import Dataset from transformers import ( AutoModelForMaskedLM, AutoTokenizer, DataCollatorForLanguageModeling, EarlyStoppingCallback, Trainer, TrainerCallback, TrainingArguments, ) class DataCollatorForWholeWordMasking(DataCollatorForLanguageModeling): """ Whole-word masking collator. When a token is selected for masking, all sub-tokens belonging to the same word are masked together. Falls back to standard token masking for tokenisers that don't set word_ids. """ def torch_call(self, examples): batch = self.tokenizer.pad(examples, return_tensors="pt", pad_to_multiple_of=self.pad_to_multiple_of) input_ids = batch["input_ids"].clone() labels = batch["input_ids"].clone() for i, (ids, encoding) in enumerate(zip(input_ids, examples)): word_ids = None if hasattr(encoding, "word_ids"): word_ids = encoding.word_ids() elif "word_ids" in encoding: word_ids = encoding["word_ids"] if word_ids is None: # Fallback: standard random token masking probability_matrix = torch.full(ids.shape, self.mlm_probability) special_tokens_mask = self.tokenizer.get_special_tokens_mask( ids.tolist(), already_has_special_tokens=True) probability_matrix[torch.tensor(special_tokens_mask, dtype=torch.bool)] = 0.0 masked_indices = torch.bernoulli(probability_matrix).bool() labels[i][~masked_indices] = -100 input_ids[i][masked_indices] = self.tokenizer.mask_token_id continue # Group token indices by word word_to_tokens: dict = {} for tok_idx, word_idx in enumerate(word_ids): if word_idx is None: continue word_to_tokens.setdefault(word_idx, []).append(tok_idx) unique_words = list(word_to_tokens.keys()) num_to_mask = max(1, int(round(len(unique_words) * self.mlm_probability))) words_to_mask = np.random.choice(unique_words, size=num_to_mask, replace=False) masked_indices = torch.zeros(ids.shape, dtype=torch.bool) for w in words_to_mask: for tok_idx in word_to_tokens[w]: if tok_idx < len(masked_indices): masked_indices[tok_idx] = True labels[i][~masked_indices] = -100 # 80% MASK, 10% random token, 10% unchanged replace_with_mask = torch.bernoulli(torch.full(ids.shape, 0.8)).bool() & masked_indices replace_with_random = (torch.bernoulli(torch.full(ids.shape, 0.5)).bool() & masked_indices & ~replace_with_mask) input_ids[i][replace_with_mask] = self.tokenizer.mask_token_id input_ids[i][replace_with_random] = torch.randint( len(self.tokenizer), ids.shape, dtype=torch.long)[replace_with_random] batch["input_ids"] = input_ids batch["labels"] = labels return batch warnings.filterwarnings("ignore") # ============================================================================= # CLEAN LOGGER (mirrors ensemble_distillation_generator.py) # ============================================================================= class LoggerTee(object): def __init__(self, filename, mode="a"): self.terminal = sys.stdout self.log = open(filename, mode) self.line_buffer = "" self.ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])') self.tqdm_pattern = re.compile(r'\b\d+%\s*\|') def write(self, message): self.terminal.write(message) for char in message: if char == '\r': self.line_buffer = "" elif char == '\n': clean_line = self.ansi_escape.sub('', self.line_buffer) if self.tqdm_pattern.search(clean_line): if "100%|" not in clean_line.replace(" ", ""): self.line_buffer = "" continue self.log.write(clean_line + '\n') self.log.flush() self.line_buffer = "" else: self.line_buffer += char def flush(self): self.terminal.flush() self.log.flush() def isatty(self): return self.terminal.isatty() timestamp = datetime.now().strftime("%Y%m%d_%H%M") log_filename = f"pretrain_log_{timestamp}.txt" sys.stdout = LoggerTee(filename=log_filename) sys.stderr = sys.stdout # ============================================================================= # ARGUMENT PARSING # ============================================================================= def parse_args(): parser = argparse.ArgumentParser( description="Domain-Adaptive MLM Pretrainer for BERT-family models." ) # --- Data --- parser.add_argument("--data_path", type=str, required=True, help="Path to unlabelled corpus (.parquet or .csv).") parser.add_argument("--text_col", type=str, required=True, help="Column name containing raw text.") parser.add_argument("--eval_split", type=float, default=0.05, help="Fraction of corpus held out for perplexity evaluation.") # --- Config / output --- parser.add_argument("--pretrain_config_path", type=str, required=True, help="Path to pretrain_config.json.") parser.add_argument("--output_base_dir", type=str, default="pretrained_checkpoints", help="Root directory; each model gets its own subdirectory here.") parser.add_argument("--metadata_path", type=str, default="pretrain_metadata.json", help="JSON file recording completed models and their final perplexity.") # --- Global training defaults (all overridable per-model in JSON) --- parser.add_argument("--max_length", type=int, default=256) parser.add_argument("--batch_size", type=int, default=16) parser.add_argument("--gradient_accumulation_steps", type=int, default=1) parser.add_argument("--max_steps", type=int, default=10000) parser.add_argument("--eval_steps", type=int, default=500) parser.add_argument("--early_stopping_patience", type=int, default=5) parser.add_argument("--learning_rate", type=float, default=5e-5) parser.add_argument("--warmup_steps", type=int, default=500) parser.add_argument("--weight_decay", type=float, default=0.01) parser.add_argument("--adam_epsilon", type=float, default=1e-8) parser.add_argument("--max_grad_norm", type=float, default=1.0) parser.add_argument("--lr_scheduler_type", type=str, default="cosine", choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup", "inverse_sqrt"]) parser.add_argument("--mlm_probability", type=float, default=0.15, help="Fraction of tokens masked per sequence.") parser.add_argument("--whole_word_masking", action="store_true", help="Use whole-word masking instead of sub-token masking. " "Recommended for syllable-level tokenisers (PhoBERT, ViDeBERTa).") parser.add_argument("--freeze_layers", type=int, default=0, help="Freeze the bottom N encoder layers during pretraining.") parser.add_argument("--use_bf16", action="store_true") parser.add_argument("--seed", type=int, default=42) # --- Retrain options --- parser.add_argument("--retrain", action="store_true", help="Retrain ALL active models from their local checkpoint.") parser.add_argument("--retrain_models", type=str, nargs="+", default=None, metavar="MODEL_NAME", help="Retrain specific models by name (space-separated). " "Example: --retrain_models microsoft/mdeberta-v3-base google/rembert") return parser.parse_args() # ============================================================================= # CONFIG HELPERS — per-model overrides, same priority logic as ensemble script # ============================================================================= _PRETRAIN_ARG_SPECS = [ # Tokenisation ("max_length", "max_length", 256), # Batch / gradient ("batch_size", "batch_size", 16), ("gradient_accumulation_steps", "gradient_accumulation_steps", 1), # Optimiser ("learning_rate", "learning_rate", 5e-5), ("weight_decay", "weight_decay", 0.01), ("adam_epsilon", "adam_epsilon", 1e-8), ("max_grad_norm", "max_grad_norm", 1.0), # Schedule ("warmup_steps", "warmup_steps", 500), ("lr_scheduler_type", "lr_scheduler_type", "cosine"), # Steps / patience ("max_steps", "max_steps", 10000), ("eval_steps", "eval_steps", 500), ("early_stopping_patience", "early_stopping_patience", 5), # MLM-specific ("mlm_probability", "mlm_probability", 0.15), ("whole_word_masking", "whole_word_masking", False), # Regularisation ("freeze_layers", "freeze_layers", 0), # Precision / model flags ("use_bf16", "use_bf16", False), ("drop_token_type_ids", "drop_token_type_ids", None), # Model identity ("tokenizer_name", "tokenizer_name", None), ] def resolve_model_args(global_args, config: dict): """ Returns a namespace where every training arg is resolved with priority: per-model JSON config > global CLI args > hardcoded default """ resolved = copy.deepcopy(global_args) for json_key, attr, default in _PRETRAIN_ARG_SPECS: if json_key in config: setattr(resolved, attr, config[json_key]) elif not hasattr(resolved, attr) or getattr(resolved, attr) is None: setattr(resolved, attr, default) if not getattr(resolved, "tokenizer_name", None): resolved.tokenizer_name = config["model_name"] resolved.model_name = config["model_name"] return resolved def load_pretrain_configs(config_path: str) -> list: with open(config_path, "r") as f: return json.load(f) # ============================================================================= # ARTIFACT HELPERS # ============================================================================= def get_output_dir(output_base: str, model_name: str) -> Path: safe_name = model_name.replace("/", "__") return Path(output_base) / safe_name def pretrain_artifact_exists(output_base: str, model_name: str) -> bool: """A completed pretraining run leaves a config.json written by save_pretrained.""" return (get_output_dir(output_base, model_name) / "config.json").exists() def load_pretrain_metadata(output_base: str, model_name: str) -> dict: meta_path = get_output_dir(output_base, model_name) / "pretrain_metadata.json" if meta_path.exists(): with open(meta_path) as f: return json.load(f) return {} def save_pretrain_metadata(output_dir: Path, perplexity: float, eval_loss: float, model_name: str, args): meta = { "model_name": model_name, "perplexity": round(perplexity, 4), "eval_loss": round(eval_loss, 6), "mlm_probability": args.mlm_probability, "whole_word_masking": args.whole_word_masking, "max_steps": args.max_steps, "learning_rate": args.learning_rate, "pretrained_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), } with open(output_dir / "pretrain_metadata.json", "w") as f: json.dump(meta, f, indent=2) def mark_model_pretrained(config_path: str, model_name: str, perplexity: float, output_dir: str): """Atomically update the JSON config to mark a model as done.""" configs = load_pretrain_configs(config_path) for cfg in configs: if cfg["model_name"] == model_name: cfg["pretrained"] = True cfg["pretrained_at"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S") cfg["perplexity"] = round(perplexity, 4) cfg["output_dir"] = str(output_dir) break tmp_path = config_path + ".tmp" with open(tmp_path, "w") as f: json.dump(configs, f, indent=2) Path(tmp_path).replace(config_path) print(f" -> [Checkpoint] Marked '{model_name}' as pretrained " f"(Perplexity={perplexity:.4f})") # ============================================================================= # FREEZE HELPER (identical to ensemble script) # ============================================================================= def freeze_encoder_layers(model, num_layers: int): if num_layers <= 0: return encoder = None for attr in ["encoder", "bert", "deberta", "electra", "roberta", "rembert"]: enc = getattr(model, attr, None) if enc is None: # Some architectures nest the base model one level deeper # e.g. model.rembert.encoder enc = getattr(getattr(model, "rembert", None), attr, None) if enc is not None: encoder = getattr(enc, "layer", None) if encoder is not None: break if encoder is None: print(" [WARNING] Could not locate encoder layers to freeze — skipping.") return actual = min(num_layers, len(encoder)) for i in range(actual): for param in encoder[i].parameters(): param.requires_grad = False print(f" -> Froze {actual}/{len(encoder)} encoder layers.") # ============================================================================= # CALLBACKS # ============================================================================= class FormattedEvalCallback(TrainerCallback): """Prints a clean evaluation summary after each eval step.""" def on_evaluate(self, args, state, control, metrics=None, **kwargs): if metrics is None: return loss = metrics.get("eval_loss") ppl = metrics.get("eval_perplexity") step = state.global_step print(f"\n [Eval @ step {step}] loss: {loss:.4f} | " f"perplexity: {ppl:.4f}" if (loss and ppl) else f"\n [Eval @ step {step}] {metrics}") if state.best_metric is not None: print(f" Best so far → {state.best_metric:.4f} (eval_loss)") print("-" * 60) # ============================================================================= # METRICS # ============================================================================= def compute_mlm_metrics(eval_pred): """Perplexity from MLM eval loss. The Trainer passes (logits, labels) but for MLM we only need the scalar loss, which the Trainer already computes and logs as eval_loss. We add perplexity here as a derived metric.""" # eval_pred.predictions is logits — we derive perplexity from eval_loss # which the Trainer makes available via model output. The cleanest way: # return an empty dict here and compute perplexity in the callback using # state.log_history. However, HF Trainer also accepts returning it from # a custom compute_loss. The approach below is simpler and standard. logits = eval_pred.predictions # [N, seq_len, vocab_size] # We cannot compute loss here without labels; perplexity is instead # injected via a post-eval callback that reads eval_loss from metrics. return {} class PerplexityCallback(TrainerCallback): """Injects eval_perplexity into the metrics dict after each evaluation.""" def on_evaluate(self, args, state, control, metrics=None, **kwargs): if metrics and "eval_loss" in metrics: try: metrics["eval_perplexity"] = float( torch.exp(torch.tensor(metrics["eval_loss"])).item() ) except Exception: metrics["eval_perplexity"] = float("inf") # ============================================================================= # TOKENISATION & DATASET HELPERS # ============================================================================= def build_dataset(hf_dataset, tokenizer, max_length: int, cache_path: str, num_proc: int = 4): """ Tokenise a HuggingFace Dataset for MLM. - hf_dataset : already-split HF Dataset (Arrow-backed, memory-mapped) - cache_path : path for the tokenized Arrow cache file; reused on reruns - num_proc : parallel tokenisation workers The tokenized result is written to disk so it never fully resides in RAM. """ def tokenize_fn(examples): return tokenizer( examples["text"], truncation=True, max_length=max_length, padding=False, return_special_tokens_mask=True, ) ds = hf_dataset.map( tokenize_fn, batched=True, batch_size=1000, num_proc=num_proc, remove_columns=["text"], keep_in_memory=False, cache_file_name=cache_path, desc="Tokenising", ) return ds # ============================================================================= # CORE PRETRAINING FUNCTION # ============================================================================= def pretrain_model(args, config: dict, train_ds_raw, eval_ds_raw, cache_dir: Path): """ Domain-adaptive MLM pretraining for a single model. Saves the full ForMaskedLM checkpoint to output_dir. Downstream fine-tuning scripts can load this with AutoModelForSequenceClassification.from_pretrained( output_dir, ignore_mismatched_sizes=True) and the MLM head is simply ignored. """ model_name = config["model_name"] output_dir = get_output_dir(args.output_base_dir, model_name) output_dir.mkdir(parents=True, exist_ok=True) tmp_ckpt_dir = output_dir / "tmp_checkpoints" print(f"\n{'='*70}") print(f"PRETRAINING: {model_name}") print(f" -> output_dir: {output_dir}") print(f" -> tokenizer: {args.tokenizer_name}") print(f" -> mlm_probability: {args.mlm_probability}") print(f" -> whole_word_masking: {args.whole_word_masking}") print(f" -> max_length: {args.max_length}") print(f" -> max_steps: {args.max_steps} | eval_steps: {args.eval_steps}") print(f" -> lr: {args.learning_rate} | warmup: {args.warmup_steps} " f"| scheduler: {args.lr_scheduler_type}") print(f" -> batch_size: {args.batch_size} " f"| grad_accum: {args.gradient_accumulation_steps}") print(f" -> freeze_layers: {args.freeze_layers}") print(f"{'='*70}") # ------------------------------------------------------------------------- # Tokeniser # ------------------------------------------------------------------------- tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name) # ------------------------------------------------------------------------- # Datasets # ------------------------------------------------------------------------- safe_name = args.model_name.replace("/", "__") cache_base = cache_dir / safe_name cache_base.mkdir(parents=True, exist_ok=True) print(f"\n Tokenising {len(train_ds_raw):,} train / {len(eval_ds_raw):,} texts " f"(disk-cached at {cache_base}) ...") train_ds = build_dataset(train_ds_raw, tokenizer, args.max_length, cache_path=str(cache_base / "train.arrow")) eval_ds = build_dataset(eval_ds_raw, tokenizer, args.max_length, cache_path=str(cache_base / "eval.arrow")) print(f" Done. Train examples: {len(train_ds):,} | Eval examples: {len(eval_ds):,}") # ------------------------------------------------------------------------- # Data collator — WWM vs standard # ------------------------------------------------------------------------- if args.whole_word_masking: print(" -> Using DataCollatorForWholeWordMasking") collator = DataCollatorForWholeWordMasking( tokenizer=tokenizer, mlm=True, mlm_probability=args.mlm_probability, ) else: print(" -> Using DataCollatorForLanguageModeling (sub-token masking)") collator = DataCollatorForLanguageModeling( tokenizer=tokenizer, mlm=True, mlm_probability=args.mlm_probability, ) # ------------------------------------------------------------------------- # Model # ------------------------------------------------------------------------- # drop_token_type_ids auto-detection: RoBERTa/XLM-R/DeBERTa-v2 don't use them model_type = "" try: from transformers import AutoConfig cfg = AutoConfig.from_pretrained(model_name) model_type = getattr(cfg, "model_type", "") except Exception: pass _no_tti_types = ["xlm-roberta", "roberta", "camembert", "deberta-v2", "distilbert", "bart", "longformer"] drop_tti = (args.drop_token_type_ids if args.drop_token_type_ids is not None else model_type in _no_tti_types) if drop_tti: print(f" -> Dropping token_type_ids (model_type='{model_type}')") if "token_type_ids" in train_ds.column_names: train_ds = train_ds.remove_columns(["token_type_ids"]) if "token_type_ids" in eval_ds.column_names: eval_ds = eval_ds.remove_columns(["token_type_ids"]) model = AutoModelForMaskedLM.from_pretrained( args.load_from, # local checkpoint path or original HF model name use_safetensors=True, ignore_mismatched_sizes=True, torch_dtype=torch.float32, ).to("cuda") freeze_encoder_layers(model, args.freeze_layers) total_params = sum(p.numel() for p in model.parameters()) trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f" -> Parameters: {total_params:,} total | {trainable_params:,} trainable") # ------------------------------------------------------------------------- # Training arguments # ------------------------------------------------------------------------- training_args = TrainingArguments( output_dir=str(tmp_ckpt_dir), max_steps=args.max_steps, per_device_train_batch_size=args.batch_size, per_device_eval_batch_size=args.batch_size, gradient_accumulation_steps=args.gradient_accumulation_steps, learning_rate=args.learning_rate, weight_decay=args.weight_decay, adam_epsilon=args.adam_epsilon, max_grad_norm=args.max_grad_norm, warmup_steps=args.warmup_steps, lr_scheduler_type=args.lr_scheduler_type, eval_strategy="steps", eval_steps=args.eval_steps, save_strategy="steps", save_steps=args.eval_steps, load_best_model_at_end=True, metric_for_best_model="eval_loss", greater_is_better=False, logging_strategy="steps", logging_steps=args.eval_steps, save_total_limit=1, report_to="none", bf16=args.use_bf16, prediction_loss_only=True, # MLM: no need to return logits during eval ) # ------------------------------------------------------------------------- # Trainer # ------------------------------------------------------------------------- trainer = Trainer( model=model, args=training_args, train_dataset=train_ds, eval_dataset=eval_ds, data_collator=collator, callbacks=[ EarlyStoppingCallback(args.early_stopping_patience), PerplexityCallback(), FormattedEvalCallback(), ], ) # ------------------------------------------------------------------------- # Train # ------------------------------------------------------------------------- trainer.train() # ------------------------------------------------------------------------- # Evaluate best checkpoint # ------------------------------------------------------------------------- eval_results = trainer.evaluate() eval_loss = eval_results.get("eval_loss", float("inf")) perplexity = float(torch.exp(torch.tensor(eval_loss)).item()) print(f"\n [Final Eval] loss: {eval_loss:.4f} | perplexity: {perplexity:.4f}") # ------------------------------------------------------------------------- # Persist — save the full ForMaskedLM model + tokeniser # Downstream scripts use ignore_mismatched_sizes=True to discard the MLM head # ------------------------------------------------------------------------- trainer.save_model(str(output_dir)) tokenizer.save_pretrained(str(output_dir)) save_pretrain_metadata(output_dir, perplexity, eval_loss, model_name, args) print(f" -> Saved model + tokenizer → {output_dir}") # Clean up temp checkpoints if tmp_ckpt_dir.exists(): shutil.rmtree(tmp_ckpt_dir) del trainer, model torch.cuda.empty_cache() gc.collect() return perplexity # ============================================================================= # MAIN ORCHESTRATOR # ============================================================================= def main(): args = parse_args() torch.manual_seed(args.seed) np.random.seed(args.seed) Path(args.output_base_dir).mkdir(parents=True, exist_ok=True) # ------------------------------------------------------------------------- # Load corpus # ------------------------------------------------------------------------- print(f"\nLoading corpus from: {args.data_path}") # Load as a memory-mapped Arrow dataset — never materialises the full # text list in Python RAM. from datasets import load_dataset as hf_load_dataset ext = args.data_path.rsplit(".", 1)[-1].lower() fmt = "parquet" if ext == "parquet" else "csv" raw_ds = hf_load_dataset(fmt, data_files=args.data_path, split="train") # Drop nulls and keep only the text column raw_ds = raw_ds.select_columns([args.text_col]) raw_ds = raw_ds.filter(lambda x: x[args.text_col] is not None and str(x[args.text_col]).strip() != "", num_proc=4) # Rename to a fixed key so build_dataset always sees "text" if args.text_col != "text": raw_ds = raw_ds.rename_column(args.text_col, "text") print(f" -> {len(raw_ds):,} texts loaded (memory-mapped).") # Split — still in Arrow, no Python list created splits = raw_ds.train_test_split(test_size=args.eval_split, seed=args.seed) train_ds_raw = splits["train"] eval_ds_raw = splits["test"] # Cap eval at 10k — perplexity over 380k rows wastes 20 min per eval max_eval = 10_000 if len(eval_ds_raw) > max_eval: eval_ds_raw = eval_ds_raw.select(range(max_eval)) cache_dir = Path(args.output_base_dir) / "_token_cache" cache_dir.mkdir(parents=True, exist_ok=True) print(f" -> Train: {len(train_ds_raw):,} | Eval: {len(eval_ds_raw):,} (capped at {max_eval:,})") print(f" -> Token cache dir: {cache_dir}") # ------------------------------------------------------------------------- # Load pretrain configs # ------------------------------------------------------------------------- configs = load_pretrain_configs(args.pretrain_config_path) active_configs = [c for c in configs if c.get("use", True)] print(f"\n -> {len(active_configs)} model(s) active in config.") # ------------------------------------------------------------------------- # Interrupt handler — mirrors ensemble script # ------------------------------------------------------------------------- completed_models = [] _interrupted = {"flag": False} def _emergency_save(signum=None, frame=None): if _interrupted["flag"]: return _interrupted["flag"] = True sig_name = f"signal {signum}" if signum else "exit" print(f"\n\n[INTERRUPT] Caught {sig_name}.") if completed_models: with open(args.metadata_path, "w") as f: json.dump({"completed": completed_models}, f, indent=2) print(f" -> Metadata saved → {args.metadata_path}") print(f" -> Artifacts safe in: {args.output_base_dir}/") print(f" -> Re-run with the same command to resume.") else: print(" -> No models completed yet.") if signum is not None: sys.exit(1) signal.signal(signal.SIGINT, _emergency_save) signal.signal(signal.SIGTERM, _emergency_save) atexit.register(_emergency_save) # ------------------------------------------------------------------------- # Training loop # ------------------------------------------------------------------------- print("\n" + "="*70) print(f"INITIATING DOMAIN-ADAPTIVE MLM PRETRAINING ({len(active_configs)} models)") print(f"Output root: {args.output_base_dir}") print("="*70) for idx, config in enumerate(active_configs): model_name = config["model_name"] local_dir = get_output_dir(args.output_base_dir, model_name) local_exists = pretrain_artifact_exists(args.output_base_dir, model_name) # Determine whether this model is flagged for retraining is_retrain_all = getattr(args, "retrain", False) retrain_list = getattr(args, "retrain_models", None) or [] is_retrain_specific = model_name in retrain_list should_retrain = local_exists and (is_retrain_all or is_retrain_specific) should_skip = local_exists and not should_retrain if should_skip: meta = load_pretrain_metadata(args.output_base_dir, model_name) ppl = meta.get("perplexity", "?") saved_at = meta.get("pretrained_at", "?") print(f"\n\n{'*'*70}") print(f"SKIPPING {idx+1}/{len(active_configs)}: {model_name}") print(f" -> Already pretrained at {saved_at} | Perplexity: {ppl}") print(f"{'*'*70}") completed_models.append({"model_name": model_name, "perplexity": ppl, "output_dir": str(local_dir)}) continue # Resolve load_from: local checkpoint if retraining, HF hub otherwise current_args = resolve_model_args(args, config) if should_retrain: current_args.load_from = str(local_dir) print(f"\n\n{'*'*70}") print(f"RETRAINING MODEL {idx+1}/{len(active_configs)}: {model_name}") print(f" -> Loading from local checkpoint: {local_dir}") print(f"{'*'*70}") else: current_args.load_from = model_name print(f"\n\n{'*'*70}") print(f"PRETRAINING MODEL {idx+1}/{len(active_configs)}: {model_name}") print(f"{'*'*70}") perplexity = pretrain_model(current_args, config, train_ds_raw, eval_ds_raw, cache_dir) output_dir = get_output_dir(args.output_base_dir, model_name) completed_models.append({"model_name": model_name, "perplexity": round(perplexity, 4), "output_dir": str(output_dir)}) mark_model_pretrained(args.pretrain_config_path, model_name, perplexity, output_dir) # ------------------------------------------------------------------------- # Final summary # ------------------------------------------------------------------------- with open(args.metadata_path, "w") as f: json.dump({"completed": completed_models}, f, indent=2) print("\n" + "="*70) print("PRETRAINING COMPLETE — SUMMARY") print("="*70) print(f" {'Model':<45} {'Perplexity':>12} {'Output Dir'}") print(" " + "-" * 90) for entry in completed_models: name = entry["model_name"].split("/")[-1] print(f" {name:<45} {str(entry['perplexity']):>12} {entry['output_dir']}") print("="*70) print(f"\nUse the output_dir paths as `model_name` in your ensemble_config.json.") print(f"Metadata written → {args.metadata_path}") if __name__ == "__main__": main()