"""Frozen feature extraction wrappers for A1 baseline models.""" from __future__ import annotations from dataclasses import asdict, dataclass import json import os from pathlib import Path import re from typing import Any import warnings import numpy as np import pandas as pd def slugify_model_id(model_id: str) -> str: cleaned = re.sub(r"[^a-zA-Z0-9._-]+", "_", model_id.strip()) return cleaned.strip("_") or "unknown_model" def _is_llada_model_id(model_id: str) -> bool: return "llada" in model_id.lower() def _apply_llada_compat_patches(model_id: str, local_files_only: bool) -> None: """Apply compatibility patches for LLaDA remote-code models.""" try: import transformers.modeling_utils as modeling_utils if not hasattr(modeling_utils.PreTrainedModel, "all_tied_weights_keys"): modeling_utils.PreTrainedModel.all_tied_weights_keys = {} elif not isinstance(modeling_utils.PreTrainedModel.all_tied_weights_keys, dict): modeling_utils.PreTrainedModel.all_tied_weights_keys = {} except Exception: # Best effort; continue with standard load flow. pass try: from transformers import AutoConfig from transformers.dynamic_module_utils import get_class_from_dynamic_module config = AutoConfig.from_pretrained( model_id, trust_remote_code=True, local_files_only=local_files_only, ) auto_map = getattr(config, "auto_map", None) or {} class_ref = auto_map.get("AutoModelForCausalLM") if not class_ref: return model_cls = get_class_from_dynamic_module( class_ref, model_id, local_files_only=local_files_only, ) original_tie = getattr(model_cls, "tie_weights", None) if original_tie is None or getattr(original_tie, "_a1_llada_safe_wrapped", False): return def _safe_tie_weights(self: Any, *args: Any, **kwargs: Any) -> Any: kwargs.pop("missing_keys", None) kwargs.pop("recompute_mapping", None) try: return original_tie(self, *args, **kwargs) except TypeError as exc: if "unexpected keyword argument" in str(exc): return original_tie(self) raise _safe_tie_weights._a1_llada_safe_wrapped = True setattr(model_cls, "tie_weights", _safe_tie_weights) except Exception: # Best effort; continue with standard load flow. pass def _normalize_all_tied_weights_keys(model: Any) -> None: """Normalize missing/incompatible all_tied_weights_keys on loaded models.""" try: tied = getattr(model, "all_tied_weights_keys", None) if tied is None: model.all_tied_weights_keys = {} elif callable(tied): value = tied() model.all_tied_weights_keys = value if isinstance(value, dict) else {} elif not isinstance(tied, dict): model.all_tied_weights_keys = {} except Exception: pass def _load_causal_lm( model_id: str, model_dtype: Any, local_files_only: bool, ) -> Any: from transformers import AutoModelForCausalLM is_llada = _is_llada_model_id(model_id) if is_llada: _apply_llada_compat_patches(model_id=model_id, local_files_only=local_files_only) load_kwargs: dict[str, Any] = { "output_hidden_states": True, "dtype": model_dtype, "local_files_only": local_files_only, } if is_llada: load_kwargs["trust_remote_code"] = True try: model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs) except AttributeError as exc: if not (is_llada and "all_tied_weights_keys" in str(exc)): raise _apply_llada_compat_patches(model_id=model_id, local_files_only=local_files_only) try: model = AutoModelForCausalLM.from_pretrained( model_id, low_cpu_mem_usage=False, **load_kwargs, ) except TypeError: model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs) _normalize_all_tied_weights_keys(model) if hasattr(model, "config") and not hasattr(model.config, "use_cache"): try: model.config.use_cache = False except Exception: pass return model @dataclass(frozen=True) class FeatureExtractionRecord: """Summary row for cached run-level feature extraction.""" model_id: str model_slug: str run: int n_words: int n_layers: int hidden_dim: int unmatched_words: int max_words_per_chunk: int dry_run: bool device: str features_npz_path: str metadata_json_path: str def _parse_model_max_length(model: Any, tokenizer: Any) -> int: candidate_values: list[int] = [] max_position_embeddings = getattr(getattr(model, "config", None), "max_position_embeddings", None) if isinstance(max_position_embeddings, int) and max_position_embeddings > 0: candidate_values.append(int(max_position_embeddings)) tokenizer_max = getattr(tokenizer, "model_max_length", None) if isinstance(tokenizer_max, int) and 0 < tokenizer_max < 100000: candidate_values.append(int(tokenizer_max)) if candidate_values: return int(min(candidate_values)) return 4096 def _build_text_and_word_spans(words: list[str]) -> tuple[str, list[tuple[int, int]]]: safe_words = [str(word) for word in words] spans: list[tuple[int, int]] = [] cursor = 0 chunks: list[str] = [] for idx, word in enumerate(safe_words): start = cursor end = start + len(word) spans.append((start, end)) chunks.append(word) cursor = end if idx < len(safe_words) - 1: chunks.append(" ") cursor += 1 return "".join(chunks), spans def _map_tokens_to_words( token_offsets: list[tuple[int, int]], word_spans: list[tuple[int, int]], ) -> tuple[list[list[int]], int]: token_to_word: list[list[int]] = [[] for _ in range(len(word_spans))] valid_token_centers: list[tuple[int, float]] = [] for token_index, (token_start, token_end) in enumerate(token_offsets): if token_end <= token_start: continue valid_token_centers.append((token_index, (token_start + token_end) * 0.5)) for word_index, (word_start, word_end) in enumerate(word_spans): overlaps = token_end > word_start and token_start < word_end if overlaps: token_to_word[word_index].append(token_index) break unmatched_words = 0 if valid_token_centers: centers = np.array([center for _, center in valid_token_centers], dtype=np.float64) indices = [idx for idx, _ in valid_token_centers] for word_index, word_tokens in enumerate(token_to_word): if word_tokens: continue unmatched_words += 1 word_start, word_end = word_spans[word_index] word_center = (word_start + word_end) * 0.5 nearest_idx = int(np.argmin(np.abs(centers - word_center))) token_to_word[word_index] = [indices[nearest_idx]] else: unmatched_words = len(token_to_word) return token_to_word, unmatched_words def _extract_chunk_features( words_chunk: list[str], model: Any, tokenizer: Any, device: str, selected_layers: list[int], ) -> tuple[dict[int, np.ndarray], dict[str, Any]]: import torch if not getattr(tokenizer, "is_fast", False): raise RuntimeError( "Fast tokenizer with offset mapping is required for word-level aggregation." ) text, word_spans = _build_text_and_word_spans(words_chunk) max_length = _parse_model_max_length(model=model, tokenizer=tokenizer) encoded = tokenizer( text, return_tensors="pt", return_offsets_mapping=True, truncation=True, max_length=max_length, add_special_tokens=True, return_overflowing_tokens=True, ) input_ids = encoded["input_ids"] if input_ids.shape[0] != 1: raise RuntimeError( "Tokenizer overflow produced multiple windows. " "Decrease --max-words-per-chunk." ) offset_mapping = encoded.pop("offset_mapping")[0].cpu().numpy().tolist() model_inputs: dict[str, Any] = {} for key, value in encoded.items(): if key in {"overflow_to_sample_mapping", "num_truncated_tokens"}: continue model_inputs[key] = value.to(device) with torch.no_grad(): try: outputs = model(**model_inputs, output_hidden_states=True, use_cache=False) except TypeError as exc: if "unexpected keyword argument" not in str(exc) or "use_cache" not in str(exc): raise outputs = model(**model_inputs, output_hidden_states=True) hidden_states = outputs.hidden_states if hidden_states is None: raise RuntimeError("Model did not return hidden states") token_to_word, unmatched_words = _map_tokens_to_words( token_offsets=[(int(start), int(end)) for start, end in offset_mapping], word_spans=word_spans, ) per_layer_features: dict[int, np.ndarray] = {} hidden_dim = int(hidden_states[selected_layers[0]].shape[-1]) for layer_idx in selected_layers: layer_tokens = hidden_states[layer_idx][0].detach().float().cpu().numpy() layer_word = np.zeros((len(words_chunk), hidden_dim), dtype=np.float32) for word_index, token_indices in enumerate(token_to_word): valid = [idx for idx in token_indices if 0 <= idx < layer_tokens.shape[0]] if not valid: continue layer_word[word_index] = np.mean(layer_tokens[valid], axis=0, dtype=np.float32) per_layer_features[layer_idx] = layer_word diagnostics = { "n_words": int(len(words_chunk)), "n_tokens": int(len(offset_mapping)), "unmatched_words": int(unmatched_words), } return per_layer_features, diagnostics def _extract_real_features_for_run( words: list[str], model: Any, tokenizer: Any, device: str, layer_indices: list[int] | None, max_words_per_chunk: int, ) -> tuple[dict[int, np.ndarray], dict[str, Any]]: if max_words_per_chunk <= 0: raise ValueError("max_words_per_chunk must be positive") if not words: raise ValueError("Cannot extract features from an empty word list") n_all_layers = int(getattr(model.config, "num_hidden_layers", 0)) + 1 selected_layers = layer_indices if layer_indices is not None else list(range(n_all_layers)) for layer_idx in selected_layers: if layer_idx < 0 or layer_idx >= n_all_layers: raise ValueError(f"Layer index {layer_idx} out of range [0, {n_all_layers - 1}]") chunk_outputs: dict[int, list[np.ndarray]] = {layer_idx: [] for layer_idx in selected_layers} total_unmatched_words = 0 total_tokens = 0 start = 0 while start < len(words): stop = min(start + max_words_per_chunk, len(words)) chunk_words = words[start:stop] per_layer_chunk, chunk_diag = _extract_chunk_features( words_chunk=chunk_words, model=model, tokenizer=tokenizer, device=device, selected_layers=selected_layers, ) total_unmatched_words += int(chunk_diag["unmatched_words"]) total_tokens += int(chunk_diag["n_tokens"]) for layer_idx in selected_layers: chunk_outputs[layer_idx].append(per_layer_chunk[layer_idx]) start = stop outputs: dict[int, np.ndarray] = { layer_idx: np.concatenate(chunks, axis=0).astype(np.float32) for layer_idx, chunks in chunk_outputs.items() } hidden_dim = int(outputs[selected_layers[0]].shape[1]) diagnostics = { "n_words": int(len(words)), "n_layers": int(len(selected_layers)), "hidden_dim": hidden_dim, "unmatched_words": int(total_unmatched_words), "n_tokens_total": int(total_tokens), "selected_layers": selected_layers, } return outputs, diagnostics def _extract_dry_run_features_for_run( words: list[str], model_id: str, run: int, dry_run_n_layers: int, dry_run_hidden_dim: int, ) -> tuple[dict[int, np.ndarray], dict[str, Any]]: if dry_run_n_layers <= 0: raise ValueError("dry_run_n_layers must be positive") if dry_run_hidden_dim <= 0: raise ValueError("dry_run_hidden_dim must be positive") n_words = len(words) seed = abs(hash((model_id, int(run), n_words))) % (2**32) rng = np.random.default_rng(seed) outputs: dict[int, np.ndarray] = {} for layer_idx in range(dry_run_n_layers): features = rng.standard_normal(size=(n_words, dry_run_hidden_dim)).astype(np.float32) outputs[layer_idx] = features diagnostics = { "n_words": int(n_words), "n_layers": int(dry_run_n_layers), "hidden_dim": int(dry_run_hidden_dim), "unmatched_words": 0, "n_tokens_total": int(n_words), "selected_layers": list(range(dry_run_n_layers)), } return outputs, diagnostics def extract_and_cache_run_level_features( run_events_df: pd.DataFrame, model_ids: list[str], output_dir: Path, layer_indices: list[int] | None, max_words_per_chunk: int, dry_run: bool, dry_run_n_layers: int, dry_run_hidden_dim: int, device: str, local_files_only: bool, overwrite: bool, num_workers: int = 1, ) -> tuple[pd.DataFrame, dict[str, Any]]: """Extract and cache run-level word features for each model. When ``num_workers > 1`` and multiple CUDA devices are visible, the work is sharded across one process per GPU (each process pinned via ``CUDA_VISIBLE_DEVICES``). Runs are partitioned round-robin across workers; each worker still iterates the full ``model_ids`` list internally. """ if run_events_df.empty: raise ValueError("run_events_df is empty; cannot extract features") required_columns = {"run", "word_index", "word", "onset_s", "offset_s"} missing = required_columns.difference(run_events_df.columns) if missing: raise ValueError(f"run_events_df missing required columns: {sorted(missing)}") output_dir = output_dir.resolve() output_dir.mkdir(parents=True, exist_ok=True) runs = sorted({int(run) for run in run_events_df["run"].tolist()}) if ( num_workers > 1 and not dry_run and len(runs) > 1 and _multi_gpu_available(device=device, requested_workers=num_workers) ): return _dispatch_multi_gpu_feature_extraction( run_events_df=run_events_df, model_ids=model_ids, output_dir=output_dir, layer_indices=layer_indices, max_words_per_chunk=max_words_per_chunk, local_files_only=local_files_only, overwrite=overwrite, num_workers=num_workers, runs=runs, ) summary_rows: list[FeatureExtractionRecord] = [] for model_id in model_ids: model_slug = slugify_model_id(model_id) model_output_dir = output_dir / model_slug model_output_dir.mkdir(parents=True, exist_ok=True) resolved_device = "dry-run" model = None tokenizer = None if not dry_run: import torch from transformers import AutoTokenizer is_llada = _is_llada_model_id(model_id) if device == "auto": resolved_device = "cuda" if torch.cuda.is_available() else "cpu" else: resolved_device = device model_dtype = torch.float16 if resolved_device.startswith("cuda") else torch.float32 tokenizer = AutoTokenizer.from_pretrained( model_id, use_fast=True, local_files_only=local_files_only, trust_remote_code=is_llada, ) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token model = _load_causal_lm( model_id, model_dtype=model_dtype, local_files_only=local_files_only, ) try: model.to(resolved_device) except torch.OutOfMemoryError as exc: if device != "auto" or not resolved_device.startswith("cuda"): raise RuntimeError( "CUDA out of memory while moving model to device. " "Retry with --feature-device cpu or reduce model size." ) from exc warnings.warn( f"CUDA OOM while loading model {model_id}; falling back to CPU.", RuntimeWarning, ) try: del model torch.cuda.empty_cache() except Exception: pass resolved_device = "cpu" model_dtype = torch.float32 model = _load_causal_lm( model_id, model_dtype=model_dtype, local_files_only=local_files_only, ) model.to(resolved_device) model.eval() for run in runs: run_df = run_events_df[run_events_df["run"] == run].sort_values("word_index") words = run_df["word"].astype(str).tolist() npz_path = model_output_dir / f"run-{run:02d}_features.npz" metadata_path = model_output_dir / f"run-{run:02d}_metadata.json" if npz_path.exists() and metadata_path.exists() and not overwrite: with metadata_path.open("r", encoding="utf-8") as handle: metadata = json.load(handle) summary_rows.append( FeatureExtractionRecord( model_id=model_id, model_slug=model_slug, run=int(run), n_words=int(metadata["n_words"]), n_layers=int(metadata["n_layers"]), hidden_dim=int(metadata["hidden_dim"]), unmatched_words=int(metadata.get("unmatched_words", 0)), max_words_per_chunk=int(metadata.get("max_words_per_chunk", max_words_per_chunk)), dry_run=bool(metadata.get("dry_run", dry_run)), device=str(metadata.get("device", resolved_device)), features_npz_path=str(npz_path), metadata_json_path=str(metadata_path), ) ) continue if dry_run: feature_map, diagnostics = _extract_dry_run_features_for_run( words=words, model_id=model_id, run=int(run), dry_run_n_layers=dry_run_n_layers, dry_run_hidden_dim=dry_run_hidden_dim, ) else: assert model is not None assert tokenizer is not None try: feature_map, diagnostics = _extract_real_features_for_run( words=words, model=model, tokenizer=tokenizer, device=resolved_device, layer_indices=layer_indices, max_words_per_chunk=max_words_per_chunk, ) except torch.OutOfMemoryError as exc: if device != "auto" or not resolved_device.startswith("cuda"): raise RuntimeError( "CUDA out of memory during feature extraction. " "Retry with --feature-device cpu or reduce --max-words-per-chunk." ) from exc warnings.warn( ( f"CUDA OOM during feature extraction for model {model_id}, run={run}; " "falling back to CPU and retrying." ), RuntimeWarning, ) torch.cuda.empty_cache() resolved_device = "cpu" model.to(resolved_device) feature_map, diagnostics = _extract_real_features_for_run( words=words, model=model, tokenizer=tokenizer, device=resolved_device, layer_indices=layer_indices, max_words_per_chunk=max_words_per_chunk, ) np.savez( npz_path, **{f"layer_{layer_idx}": values for layer_idx, values in feature_map.items()}, onset_s=run_df["onset_s"].to_numpy(dtype=np.float32), offset_s=run_df["offset_s"].to_numpy(dtype=np.float32), word_index=run_df["word_index"].to_numpy(dtype=np.int64), ) metadata = { "model_id": model_id, "model_slug": model_slug, "run": int(run), "n_words": int(diagnostics["n_words"]), "n_layers": int(diagnostics["n_layers"]), "hidden_dim": int(diagnostics["hidden_dim"]), "unmatched_words": int(diagnostics["unmatched_words"]), "n_tokens_total": int(diagnostics["n_tokens_total"]), "selected_layers": [int(value) for value in diagnostics["selected_layers"]], "max_words_per_chunk": int(max_words_per_chunk), "dry_run": bool(dry_run), "device": str(resolved_device), "features_npz_path": str(npz_path), } with metadata_path.open("w", encoding="utf-8") as handle: json.dump(metadata, handle, indent=2, sort_keys=True) summary_rows.append( FeatureExtractionRecord( model_id=model_id, model_slug=model_slug, run=int(run), n_words=int(diagnostics["n_words"]), n_layers=int(diagnostics["n_layers"]), hidden_dim=int(diagnostics["hidden_dim"]), unmatched_words=int(diagnostics["unmatched_words"]), max_words_per_chunk=int(max_words_per_chunk), dry_run=bool(dry_run), device=str(resolved_device), features_npz_path=str(npz_path), metadata_json_path=str(metadata_path), ) ) if model is not None: del model del tokenizer summary_df = pd.DataFrame([asdict(row) for row in summary_rows]) if not summary_df.empty: summary_df = summary_df.sort_values(["model_slug", "run"]).reset_index(drop=True) feature_qc: dict[str, Any] = { "n_models": int(len({row.model_slug for row in summary_rows})), "n_model_run_rows": int(len(summary_rows)), "dry_run": bool(dry_run), "max_words_per_chunk": int(max_words_per_chunk), } if not summary_df.empty: feature_qc["n_layers_min"] = int(summary_df["n_layers"].min()) feature_qc["n_layers_max"] = int(summary_df["n_layers"].max()) feature_qc["hidden_dim_min"] = int(summary_df["hidden_dim"].min()) feature_qc["hidden_dim_max"] = int(summary_df["hidden_dim"].max()) feature_qc["unmatched_words_total"] = int(summary_df["unmatched_words"].sum()) return summary_df, feature_qc def _multi_gpu_available(device: str, requested_workers: int) -> bool: """Return True when CUDA exposes >=2 devices and the request is sane.""" if requested_workers <= 1: return False device_lower = str(device).strip().lower() if device_lower in {"cpu", "dry-run"}: return False try: import torch except Exception: return False if not torch.cuda.is_available(): return False return torch.cuda.device_count() >= 2 def resolve_feature_num_workers(requested: int | str, device: str) -> int: """Resolve --feature-num-workers ('auto' or int) to a concrete worker count. Returns 1 unless multiple CUDA devices are visible and ``device`` is auto/cuda. """ device_lower = str(device).strip().lower() if device_lower in {"cpu", "dry-run"}: return 1 try: import torch n_gpus = int(torch.cuda.device_count()) if torch.cuda.is_available() else 0 except Exception: n_gpus = 0 if isinstance(requested, str): token = requested.strip().lower() if token in {"", "auto"}: return max(1, n_gpus) try: value = int(token) except ValueError as exc: raise ValueError(f"Invalid --feature-num-workers={requested!r}") from exc else: value = int(requested) if value <= 1: return 1 if n_gpus <= 0: return 1 return min(value, n_gpus) def _record_from_metadata( *, model_id: str, model_slug: str, run: int, metadata: dict[str, Any], npz_path: Path, metadata_path: Path, max_words_per_chunk: int, ) -> FeatureExtractionRecord: return FeatureExtractionRecord( model_id=model_id, model_slug=model_slug, run=int(run), n_words=int(metadata["n_words"]), n_layers=int(metadata["n_layers"]), hidden_dim=int(metadata["hidden_dim"]), unmatched_words=int(metadata.get("unmatched_words", 0)), max_words_per_chunk=int(metadata.get("max_words_per_chunk", max_words_per_chunk)), dry_run=bool(metadata.get("dry_run", False)), device=str(metadata.get("device", "cuda")), features_npz_path=str(npz_path), metadata_json_path=str(metadata_path), ) def _gpu_worker_entrypoint( rank: int, world_size: int, payload_path: str, ) -> None: """Process entrypoint for one GPU worker. Runs in a spawned subprocess.""" import pickle # Pin this process to a single GPU before importing torch in the child. os.environ["CUDA_VISIBLE_DEVICES"] = str(rank) # Avoid BLAS thread oversubscription across workers. n_cpu = os.cpu_count() or 8 threads = max(1, n_cpu // max(1, world_size)) os.environ.setdefault("OMP_NUM_THREADS", str(threads)) os.environ.setdefault("MKL_NUM_THREADS", str(threads)) os.environ.setdefault("OPENBLAS_NUM_THREADS", str(threads)) os.environ.setdefault("NUMEXPR_NUM_THREADS", str(threads)) try: import torch torch.set_num_threads(threads) except Exception: pass with open(payload_path, "rb") as handle: payload: dict[str, Any] = pickle.load(handle) all_runs: list[int] = payload["runs"] my_runs = [r for idx, r in enumerate(all_runs) if idx % world_size == rank] if not my_runs: return run_events_df: pd.DataFrame = payload["run_events_df"] df_subset = run_events_df[run_events_df["run"].isin(my_runs)].reset_index(drop=True) if df_subset.empty: return extract_and_cache_run_level_features( run_events_df=df_subset, model_ids=payload["model_ids"], output_dir=Path(payload["output_dir"]), layer_indices=payload["layer_indices"], max_words_per_chunk=payload["max_words_per_chunk"], dry_run=False, dry_run_n_layers=0, dry_run_hidden_dim=0, device="cuda:0", local_files_only=payload["local_files_only"], overwrite=payload["overwrite"], num_workers=1, ) def _dispatch_multi_gpu_feature_extraction( *, run_events_df: pd.DataFrame, model_ids: list[str], output_dir: Path, layer_indices: list[int] | None, max_words_per_chunk: int, local_files_only: bool, overwrite: bool, num_workers: int, runs: list[int], ) -> tuple[pd.DataFrame, dict[str, Any]]: """Spawn one worker per GPU; each worker handles a disjoint subset of runs.""" import pickle import tempfile import torch.multiprocessing as mp world_size = min(int(num_workers), len(runs)) print( f"[features] Multi-GPU feature extraction: world_size={world_size}, " f"runs={runs}, models={len(model_ids)}", flush=True, ) payload = { "runs": runs, "run_events_df": run_events_df, "model_ids": list(model_ids), "output_dir": str(output_dir), "layer_indices": layer_indices, "max_words_per_chunk": int(max_words_per_chunk), "local_files_only": bool(local_files_only), "overwrite": bool(overwrite), } with tempfile.NamedTemporaryFile( mode="wb", suffix=".pkl", delete=False, dir=str(output_dir) ) as handle: pickle.dump(payload, handle) payload_path = handle.name try: mp.spawn( _gpu_worker_entrypoint, args=(world_size, payload_path), nprocs=world_size, join=True, ) finally: try: os.unlink(payload_path) except OSError: pass # Aggregate summary by reading metadata files written by workers. summary_rows: list[FeatureExtractionRecord] = [] for model_id in model_ids: model_slug = slugify_model_id(model_id) model_output_dir = output_dir / model_slug for run in runs: npz_path = model_output_dir / f"run-{run:02d}_features.npz" metadata_path = model_output_dir / f"run-{run:02d}_metadata.json" if not (npz_path.exists() and metadata_path.exists()): raise RuntimeError( f"Multi-GPU worker did not produce features for " f"model={model_id} run={run}: missing {metadata_path} or {npz_path}" ) with metadata_path.open("r", encoding="utf-8") as handle: metadata = json.load(handle) summary_rows.append( _record_from_metadata( model_id=model_id, model_slug=model_slug, run=int(run), metadata=metadata, npz_path=npz_path, metadata_path=metadata_path, max_words_per_chunk=int(max_words_per_chunk), ) ) summary_df = pd.DataFrame([asdict(row) for row in summary_rows]) if not summary_df.empty: summary_df = summary_df.sort_values(["model_slug", "run"]).reset_index(drop=True) feature_qc: dict[str, Any] = { "n_models": int(len({row.model_slug for row in summary_rows})), "n_model_run_rows": int(len(summary_rows)), "dry_run": False, "max_words_per_chunk": int(max_words_per_chunk), "multi_gpu_world_size": int(world_size), } if not summary_df.empty: feature_qc["n_layers_min"] = int(summary_df["n_layers"].min()) feature_qc["n_layers_max"] = int(summary_df["n_layers"].max()) feature_qc["hidden_dim_min"] = int(summary_df["hidden_dim"].min()) feature_qc["hidden_dim_max"] = int(summary_df["hidden_dim"].max()) feature_qc["unmatched_words_total"] = int(summary_df["unmatched_words"].sum()) return summary_df, feature_qc