""" MLM evaluation for safetensors-format bareminn encoder weights (fp16). Loads a model.safetensors file directly, converts weights to fp32 for bit-exact evaluation, and runs the full eval suite to prove the fp16- stored weights produce identical results to the original fp32 checkpoint. Usage: python scripts/mlm_eval_safetensors.py --repo-dir . --device cpu python scripts/mlm_eval_safetensors.py --repo-dir . --device cuda --compare eval_results/d12_encoder_mlm_eval.json """ import os import json import argparse from collections import Counter import math import torch import torch.nn.functional as F from transformers import BertTokenizer from pathlib import Path from transformers import AutoModelForMaskedLM MAX_LENGTH = 104 # --------------------------------------------------------------------------- # Device resolution # --------------------------------------------------------------------------- def resolve_device(requested: str) -> str: if requested == "cuda" and not torch.cuda.is_available(): raise RuntimeError( "Requested --device cuda, but torch.cuda.is_available() is False." ) return requested # --------------------------------------------------------------------------- # Safetensors loader (replaces load_model for .pt checkpoints) # --------------------------------------------------------------------------- def load_model_from_repo(repo_dir: Path, device: str): """Load encoder weights from a Hugging Face repo directory (config.json + model.safetensors + tokenizer files). Weights on disk are fp16; we load at fp32 so the forward pass is numerically identical to the original fp32 checkpoint — proving the fp16→fp32 round-trip is lossless for these values. """ print(f" Loading from repo: {repo_dir}") model = AutoModelForMaskedLM.from_pretrained( str(repo_dir), dtype=torch.float32, trust_remote_code=True, ) model.to(device) model.eval() return model, model.config # --------------------------------------------------------------------------- # Tokenization helpers (unchanged from original) # --------------------------------------------------------------------------- def tokenize_with_mask(text: str, tokenizer: BertTokenizer) -> list[int]: lowered = "[MASK]".join(part.lower() for part in text.split("[MASK]")) enc = tokenizer(lowered, add_special_tokens=False) return enc["input_ids"][:MAX_LENGTH] def word_to_token_ids(word: str, tokenizer: BertTokenizer) -> list[int]: enc = tokenizer.encode(word.lower(), add_special_tokens=True) return enc[1:-1] def word_to_first_token_id(word: str, tokenizer: BertTokenizer) -> int: inner = word_to_token_ids(word, tokenizer) return inner[0] if inner else tokenizer.unk_token_id def find_mask_positions(token_ids: list[int], mask_id: int) -> list[int]: return [i for i, t in enumerate(token_ids) if t == mask_id] # --------------------------------------------------------------------------- # Candidate / sequence helpers (unchanged from original) # --------------------------------------------------------------------------- def _build_candidates(tc: dict) -> list[str]: return ( tc.get("targets", []) + tc.get("foils", []) + tc.get("acceptable_alternatives", []) + tc.get("failure_examples", []) ) def _required_mask_count(tc: dict, tokenizer: BertTokenizer) -> int: existing_masks = tc["input"].count("[MASK]") targets = tc.get("targets", []) required = max((len(word_to_token_ids(w, tokenizer)) for w in targets), default=1) return max(existing_masks, required, 1) def _expand_masks(text: str, required_masks: int) -> str: existing_masks = text.count("[MASK]") if existing_masks != 1 or required_masks <= 1: return text return text.replace("[MASK]", " ".join(["[MASK]"] * required_masks), 1) def _sequence_candidates( tc: dict, tokenizer: BertTokenizer, include_failures: bool = False, ) -> list[tuple[str, tuple[int, ...]]]: seen: set[tuple[int, ...]] = set() words = tc.get("targets", []) + tc.get("acceptable_alternatives", []) if tc["pass_condition"] == "correct_beats_foil": words += tc.get("foils", []) if include_failures: words += tc.get("failure_examples", []) out = [] for word in words: ids = tuple(word_to_token_ids(word, tokenizer)) if not ids or ids in seen: continue seen.add(ids) out.append((word, ids)) return out def _sequence_probability( probs_i: torch.Tensor, mask_pos: list[int], token_ids: list[int] | tuple[int, ...], ) -> float: if not token_ids or len(token_ids) > len(mask_pos): return 0.0 log_probs = [] for pos, token_id in zip(mask_pos, token_ids): prob = probs_i[pos, token_id].item() log_probs.append(math.log(max(prob, 1e-45))) return math.exp(sum(log_probs) / len(log_probs)) def _best_windowed_sequence_probability( probs_i: torch.Tensor, mask_pos: list[int], token_ids: list[int] | tuple[int, ...], ) -> float: if not token_ids or len(token_ids) > len(mask_pos): return 0.0 best = 0.0 width = len(token_ids) for start in range(0, len(mask_pos) - width + 1): prob = _sequence_probability( probs_i, mask_pos[start : start + width], token_ids ) best = max(best, prob) return best def _candidate_sequences( words: list[str], tokenizer: BertTokenizer ) -> list[tuple[str, tuple[int, ...]]]: out = [] for word in words: ids = tuple(word_to_token_ids(word, tokenizer)) if ids: out.append((word, ids)) return out def _contains_sequence(seq: tuple[int, ...], candidate: tuple[int, ...]) -> bool: if not candidate or len(candidate) > len(seq): return False width = len(candidate) return any( seq[start : start + width] == candidate for start in range(0, len(seq) - width + 1) ) def _match_sequence( seq: tuple[int, ...], targets: list[tuple[str, tuple[int, ...]]], alternatives: list[tuple[str, tuple[int, ...]]], ) -> tuple[str, str | None]: for word, ids in targets: if seq == ids: return "exact_target", word for word, ids in alternatives: if seq == ids: return "exact_alt", word for word, ids in alternatives: if len(ids) < len(seq) and _contains_sequence(seq, ids): return "contained_alt", word return "none", None def _matches_any_sequence( seq: tuple[int, ...], candidates: list[tuple[str, tuple[int, ...]]], ) -> bool: return any( seq == ids or (len(ids) < len(seq) and _contains_sequence(seq, ids)) for _word, ids in candidates ) def _best_candidate_probability( probs_i: torch.Tensor, mask_pos: list[int], candidates: list[tuple[str, tuple[int, ...]]], ) -> tuple[float, str | None]: best_prob = 0.0 best_word = None for word, ids in candidates: prob = _best_windowed_sequence_probability(probs_i, mask_pos, ids) if prob > best_prob: best_prob = prob best_word = word return best_prob, best_word def _token_candidates(mask_probs: torch.Tensor, k: int, n_samples: int) -> list[int]: if n_samples <= 1: return torch.topk(mask_probs, k).indices.tolist() sampled: set[int] = set() draw_k = min(k, mask_probs.numel()) for _ in range(n_samples): drawn = torch.multinomial(mask_probs, draw_k, replacement=False) sampled.update(drawn.tolist()) return sorted(sampled, key=lambda i: mask_probs[i].item(), reverse=True) def _decode_token_sequence( token_ids: list[int] | tuple[int, ...], tokenizer: BertTokenizer ) -> str: return tokenizer.convert_tokens_to_string( tokenizer.convert_ids_to_tokens(list(token_ids)) ).strip() def _top_k_sequences( probs_i: torch.Tensor, mask_pos: list[int], k: int, tokenizer: BertTokenizer, n_samples: int, beam_width: int | None = None, ) -> list[tuple[tuple[int, ...], float]]: if not mask_pos: return [] beam_width = beam_width or max(20, k * 4) candidates_per_mask = [ _token_candidates(probs_i[pos], beam_width, n_samples) for pos in mask_pos ] beams: list[tuple[tuple[int, ...], float]] = [(tuple(), 0.0)] for pos, token_ids in zip(mask_pos, candidates_per_mask): next_beams: list[tuple[tuple[int, ...], float]] = [] for prefix, score in beams: for token_id in token_ids: prob = probs_i[pos, token_id].item() next_beams.append( (prefix + (token_id,), score + math.log(max(prob, 1e-45))) ) next_beams.sort(key=lambda item: item[1], reverse=True) beams = next_beams[:beam_width] return beams[:k] def _classify_error( tc: dict, mask_probs: torch.Tensor, top_k_ids: list[int], failure_ids: set, k: int, passed: bool, tokenizer: BertTokenizer, ) -> str | None: if passed: return None targets = tc.get("targets", []) if not targets: return "total_miss" primary_correct_id = word_to_first_token_id(targets[0], tokenizer) sorted_ids = torch.argsort(mask_probs, descending=True).tolist() correct_rank = sorted_ids.index(primary_correct_id) if k <= correct_rank <= k + 1: return "near_miss" if failure_ids and any(idx in failure_ids for idx in top_k_ids[:3]): return "generic_over_theological" if correct_rank >= 20: return "total_miss" return "wrong_semantic_cluster" # --------------------------------------------------------------------------- # Batching and inference (unchanged from original) # --------------------------------------------------------------------------- def prepare_batch( test_cases: list[dict], tokenizer: BertTokenizer, device: str ): mask_id = tokenizer.mask_token_id pad_id = tokenizer.pad_token_id tokenized = [] mask_positions_list = [] multi_piece_flags = [] for tc in test_cases: required_masks = _required_mask_count(tc, tokenizer) original_masks = tc["input"].count("[MASK]") expanded = _expand_masks(tc["input"], required_masks) ids = tokenize_with_mask(expanded, tokenizer) mpos = find_mask_positions(ids, mask_id) tokenized.append(ids) mask_positions_list.append(mpos) multi_piece_flags.append(original_masks == 1 and required_masks > 1) seq_len = max(len(ids) for ids in tokenized) padded = [ids + [pad_id] * (seq_len - len(ids)) for ids in tokenized] input_tensor = torch.tensor(padded, dtype=torch.long, device=device) attention_mask = (input_tensor != pad_id).long() return input_tensor, attention_mask, mask_positions_list, multi_piece_flags @torch.no_grad() def run_inference( model, input_tensor: torch.Tensor, attention_mask: torch.Tensor, temperature: float = 1.0, ) -> torch.Tensor: outputs = model(input_tensor, attention_mask=attention_mask) logits = outputs.logits if temperature != 1.0: logits = logits / temperature return F.softmax(logits, dim=-1) def _predicted_tokens( mask_probs: torch.Tensor, k: int, n_samples: int, tokenizer: BertTokenizer, ) -> list[str]: ids = _token_candidates(mask_probs, k, n_samples) return tokenizer.convert_ids_to_tokens(ids) # --------------------------------------------------------------------------- # Scoring functions (unchanged from original) # --------------------------------------------------------------------------- def score_target_in_top_k( tc, probs_i, mask_pos, k, n_samples, tokenizer, multi_piece_target=False ): if multi_piece_target: top_sequences = _top_k_sequences( probs_i, mask_pos, k, tokenizer, n_samples ) top_k_ids = [seq for seq, _ in top_sequences] top_k_tokens = [ _decode_token_sequence(seq, tokenizer) for seq, _ in top_sequences ] targets = tc.get("targets", []) alts = tc.get("acceptable_alternatives", []) target_ids = _candidate_sequences(targets, tokenizer) alt_ids = _candidate_sequences(alts, tokenizer) passed = False mrr = 0.0 match_type = "none" matched_candidate = None for rank, seq in enumerate(top_k_ids): match_type, matched_candidate = _match_sequence( seq, target_ids, alt_ids ) if match_type != "none": passed = True mrr = 1.0 / (rank + 1) break failure_ids = _candidate_sequences( tc.get("failure_examples", []), tokenizer ) critical_failure = any( _matches_any_sequence(seq, failure_ids) for seq in top_k_ids[:3] ) target_rank = None for _word, ids in target_ids: for rank, beam in enumerate(top_k_ids): if beam == ids: target_rank = rank break if target_rank is not None: break if passed: error_type = None elif target_rank is not None and k <= target_rank <= k + 1: error_type = "near_miss" elif critical_failure: error_type = "generic_over_theological" else: error_type = "total_miss" return dict( id=tc["id"], type=tc["type"], category=tc["category"], difficulty=tc["difficulty"], pass_=passed, mrr=mrr, margin=None, confidence=None, critical_failure=critical_failure, masked_sentence=tc["input"], correct_token=targets[0] if targets else "", reference=tc.get("reference", ""), candidates=_build_candidates(tc), predicted_top_k=top_k_tokens, error_type=error_type, match_type=match_type, matched_candidate=matched_candidate, ) first_mask = mask_pos[0] if mask_pos else 0 mask_probs = probs_i[first_mask] top_k_ids = torch.topk(mask_probs, k).indices.tolist() top_k_tokens = _predicted_tokens(mask_probs, k, n_samples, tokenizer) targets = tc.get("targets", []) alts = tc.get("acceptable_alternatives", []) correct_ids = set( word_to_first_token_id(w, tokenizer) for w in targets + alts ) passed = any(idx in correct_ids for idx in top_k_ids) mrr = 0.0 for rank, idx in enumerate(top_k_ids): if idx in correct_ids: mrr = 1.0 / (rank + 1) break failure_ids = set( word_to_first_token_id(w, tokenizer) for w in tc.get("failure_examples", []) ) critical_failure = any(idx in failure_ids for idx in top_k_ids[:3]) error_type = _classify_error( tc, mask_probs, top_k_ids, failure_ids, k, passed, tokenizer ) return dict( id=tc["id"], type=tc["type"], category=tc["category"], difficulty=tc["difficulty"], pass_=passed, mrr=mrr, margin=None, confidence=None, critical_failure=critical_failure, masked_sentence=tc["input"], correct_token=targets[0] if targets else "", reference=tc.get("reference", ""), candidates=_build_candidates(tc), predicted_top_k=top_k_tokens, error_type=error_type, ) def score_all_top_k_in_target_set( tc, probs_i, mask_pos, k, n_samples, tokenizer, multi_piece_target=False ): if multi_piece_target: top_sequences = _top_k_sequences( probs_i, mask_pos, k, tokenizer, n_samples ) top_k_ids = [seq for seq, _ in top_sequences] top_k_tokens = [ _decode_token_sequence(seq, tokenizer) for seq, _ in top_sequences ] targets = tc.get("targets", []) alts = tc.get("acceptable_alternatives", []) target_ids = _candidate_sequences(targets, tokenizer) alt_ids = _candidate_sequences(alts, tokenizer) matches = [ _match_sequence(seq, target_ids, alt_ids) for seq in top_k_ids ] valid_count = sum( 1 for match_type, _word in matches if match_type != "none" ) precision = valid_count / k passed = precision >= 0.8 failure_ids = _candidate_sequences( tc.get("failure_examples", []), tokenizer ) critical_failure = any( _matches_any_sequence(seq, failure_ids) for seq in top_k_ids ) error_type = ( None if passed else ( "generic_over_theological" if critical_failure else "wrong_semantic_cluster" ) ) matched_candidates = [ word for match_type, word in matches if match_type != "none" and word is not None ] return dict( id=tc["id"], type=tc["type"], category=tc["category"], difficulty=tc["difficulty"], pass_=passed, mrr=None, margin=precision, confidence=None, critical_failure=critical_failure, masked_sentence=tc["input"], correct_token=targets[0] if targets else "", reference=tc.get("reference", ""), candidates=_build_candidates(tc), predicted_top_k=top_k_tokens, error_type=error_type, match_type="valid_set" if passed else "none", matched_candidate=", ".join(matched_candidates) if matched_candidates else None, ) first_mask = mask_pos[0] if mask_pos else 0 mask_probs = probs_i[first_mask] top_k_ids = torch.topk(mask_probs, k).indices.tolist() top_k_tokens = _predicted_tokens(mask_probs, k, n_samples, tokenizer) targets = tc.get("targets", []) alts = tc.get("acceptable_alternatives", []) valid_ids = set( word_to_first_token_id(w, tokenizer) for w in targets + alts ) valid_count = sum(1 for idx in top_k_ids if idx in valid_ids) precision = valid_count / k passed = precision >= 0.8 failure_ids = set( word_to_first_token_id(w, tokenizer) for w in tc.get("failure_examples", []) ) critical_failure = any(idx in failure_ids for idx in top_k_ids) error_type = _classify_error( tc, mask_probs, top_k_ids, failure_ids, k, passed, tokenizer ) return dict( id=tc["id"], type=tc["type"], category=tc["category"], difficulty=tc["difficulty"], pass_=passed, mrr=None, margin=precision, confidence=None, critical_failure=critical_failure, masked_sentence=tc["input"], correct_token=targets[0] if targets else "", reference=tc.get("reference", ""), candidates=_build_candidates(tc), predicted_top_k=top_k_tokens, error_type=error_type, ) def score_correct_beats_foil( tc, probs_i, mask_pos, k, n_samples, tokenizer, multi_piece_target=False ): if multi_piece_target: targets = tc.get("targets", []) alts = tc.get("acceptable_alternatives", []) foils = tc.get("foils", []) target_ids = _candidate_sequences(targets, tokenizer) alt_ids = _candidate_sequences(alts, tokenizer) p_correct, matched_candidate = _best_candidate_probability( probs_i, mask_pos, target_ids + alt_ids ) match_type = "none" if matched_candidate is not None: match_type = ( "exact_target" if matched_candidate in targets else "acceptable_alt" ) foil_ids = word_to_token_ids(foils[0], tokenizer) if foils else [] p_foil = ( _best_windowed_sequence_probability(probs_i, mask_pos, foil_ids) if foil_ids else 0.0 ) passed = p_correct > p_foil margin = p_correct - p_foil if margin > 0.1: confidence = "high" elif margin > 0.02: confidence = "medium" else: confidence = "low" top_sequences = _top_k_sequences( probs_i, mask_pos, k, tokenizer, n_samples ) top_k_ids = [seq for seq, _ in top_sequences] top_k_tokens = [ _decode_token_sequence(seq, tokenizer) for seq, _ in top_sequences ] failure_ids: set = set() error_type = None if passed else "wrong_semantic_cluster" return dict( id=tc["id"], type=tc["type"], category=tc["category"], difficulty=tc["difficulty"], pass_=passed, mrr=None, margin=margin, confidence=confidence, critical_failure=False, masked_sentence=tc["input"], correct_token=targets[0] if targets else "", reference=tc.get("reference", ""), candidates=_build_candidates(tc), predicted_top_k=top_k_tokens, error_type=error_type, match_type=match_type if passed else "none", matched_candidate=matched_candidate if passed else None, ) first_mask = mask_pos[0] if mask_pos else 0 mask_probs = probs_i[first_mask] targets = tc.get("targets", []) alts = tc.get("acceptable_alternatives", []) foils = tc.get("foils", []) correct_ids = [ word_to_first_token_id(w, tokenizer) for w in targets + alts ] p_correct = ( max(mask_probs[i].item() for i in correct_ids) if correct_ids else 0.0 ) foil_id = word_to_first_token_id(foils[0], tokenizer) if foils else None p_foil = mask_probs[foil_id].item() if foil_id is not None else 0.0 passed = p_correct > p_foil margin = p_correct - p_foil if margin > 0.1: confidence = "high" elif margin > 0.02: confidence = "medium" else: confidence = "low" top_k_ids = torch.topk(mask_probs, k).indices.tolist() top_k_tokens = _predicted_tokens(mask_probs, k, n_samples, tokenizer) failure_ids: set = set() error_type = _classify_error( tc, mask_probs, top_k_ids, failure_ids, k, passed, tokenizer ) return dict( id=tc["id"], type=tc["type"], category=tc["category"], difficulty=tc["difficulty"], pass_=passed, mrr=None, margin=margin, confidence=confidence, critical_failure=False, masked_sentence=tc["input"], correct_token=targets[0] if targets else "", reference=tc.get("reference", ""), candidates=_build_candidates(tc), predicted_top_k=top_k_tokens, error_type=error_type, ) def score_all( test_cases, probs, mask_positions_list, multi_piece_flags, tokenizer, default_k, n_samples=1, ): results = [] for i, tc in enumerate(test_cases): mpos = mask_positions_list[i] multi_piece_target = multi_piece_flags[i] k = tc.get("k") or default_k cond = tc["pass_condition"] probs_i = probs[i] if cond == "target_in_top_k": result = score_target_in_top_k( tc, probs_i, mpos, k, n_samples, tokenizer, multi_piece_target=multi_piece_target, ) elif cond == "all_top_k_in_target_set": result = score_all_top_k_in_target_set( tc, probs_i, mpos, k, n_samples, tokenizer, multi_piece_target=multi_piece_target, ) elif cond == "correct_beats_foil": result = score_correct_beats_foil( tc, probs_i, mpos, k, n_samples, tokenizer, multi_piece_target=multi_piece_target, ) else: raise ValueError(f"Unknown pass_condition: {cond}") results.append(result) return results # --------------------------------------------------------------------------- # Reporting (unchanged from original) # --------------------------------------------------------------------------- def pct(val): return f"{val * 100:.1f}%" def difficulty_weighted_score(results): weights = {"easy": 1.0, "medium": 2.0, "hard": 3.0} ws = wt = 0.0 for result in results: weight = weights.get(result["difficulty"], 1.0) ws += weight if result["pass_"] else 0.0 wt += weight return ws / wt if wt > 0 else 0.0 def print_results(results): total = len(results) passed = sum(1 for r in results if r["pass_"]) overall = passed / total print("\n" + "=" * 70) print("MLM EVALUATION RESULTS (safetensors fp16 → fp32)") print("=" * 70) print(f"\nOverall: {passed}/{total} ({pct(overall)})") print("\n--- BY TYPE ---") for eval_type in [ "doctrinal_association", "canonical_knowledge", "contrastive_theology", ]: matches = [r for r in results if r["type"] == eval_type] passed_matches = sum(1 for r in matches if r["pass_"]) rate = passed_matches / len(matches) if matches else 0.0 extra = "" if eval_type == "doctrinal_association": mrrs = [r["mrr"] for r in matches if r["mrr"] is not None] extra = ( f" mean_mrr={sum(mrrs) / len(mrrs):.3f}" if mrrs else "" ) elif eval_type == "contrastive_theology": margins = [r["margin"] for r in matches if r["margin"] is not None] extra = ( f" mean_margin={sum(margins) / len(margins):.4f}" if margins else "" ) print( f" {eval_type}: {passed_matches}/{len(matches)} ({pct(rate)}){extra}" ) print("\n--- BY CATEGORY ---") categories = sorted(set(r["category"] for r in results)) for category in categories: matches = [r for r in results if r["category"] == category] passed_matches = sum(1 for r in matches if r["pass_"]) rate = passed_matches / len(matches) if matches else 0.0 critical = sum(1 for r in matches if r["critical_failure"]) print( f" {category:<22s} {passed_matches}/{len(matches)} ({pct(rate):<6s}) crit_fail={critical}" ) print("\n--- BY DIFFICULTY ---") for difficulty in ["easy", "medium", "hard"]: matches = [r for r in results if r["difficulty"] == difficulty] passed_matches = sum(1 for r in matches if r["pass_"]) rate = passed_matches / len(matches) if matches else 0.0 print(f" {difficulty:<8} {passed_matches}/{len(matches)} ({pct(rate)})") weighted = difficulty_weighted_score(results) critical = sum(1 for r in results if r["critical_failure"]) print(f"\nDifficulty-weighted score: {pct(weighted)}") print( f"Critical failure rate: {critical}/{total} ({pct(critical / total)})" ) contrastive = [r for r in results if r["type"] == "contrastive_theology"] if contrastive: print("\n--- CONTRASTIVE CONFIDENCE ---") for confidence in ["high", "medium", "low"]: matches = [r for r in contrastive if r["confidence"] == confidence] passed_matches = sum(1 for r in matches if r["pass_"]) print(f" {confidence:<8s} {passed_matches}/{len(matches)}") easy_results = [r for r in results if r["difficulty"] == "easy"] easy_rate = ( sum(1 for r in easy_results if r["pass_"]) / len(easy_results) if easy_results else 0.0 ) contrastive_rate = ( sum(1 for r in contrastive if r["pass_"]) / len(contrastive) if contrastive else 0.0 ) if overall >= 0.80 and easy_rate >= 0.95 and contrastive_rate >= 0.75: assessment = "EXCELLENT" elif overall >= 0.60 and easy_rate >= 0.85 and contrastive_rate >= 0.60: assessment = "GOOD ENOUGH for semantic search" elif overall >= 0.30: assessment = "NEEDS MORE TRAINING" else: assessment = "FUNDAMENTALLY BROKEN" print("\n" + "=" * 70) print(f"ASSESSMENT: {assessment}") print("=" * 70 + "\n") # --------------------------------------------------------------------------- # Comparison against a reference eval JSON # --------------------------------------------------------------------------- def compare_results( results: list[dict], reference_path: str, ) -> dict: """Diff the safetensors eval against a prior (fp32 .pt checkpoint) run. Returns a comparison summary. Prints any test cases whose pass/fail status differs between the two runs. """ with open(reference_path) as f: ref = json.load(f) ref_results = ref.get("results", ref) # Build lookup by test case id ref_by_id: dict[str, dict] = {} for r in ref_results: rid = str(r.get("id", "")) if rid: ref_by_id[rid] = r total = len(results) matching = 0 mismatches: list[dict] = [] only_new = [] only_old_ids = set(ref_by_id.keys()) for r in results: rid = str(r.get("id", "")) only_old_ids.discard(rid) ref_r = ref_by_id.get(rid) if ref_r is None: only_new.append(r) continue pass_match = r["pass_"] == ref_r.get("pass_") if pass_match: matching += 1 else: mismatches.append( { "id": rid, "type": r["type"], "category": r["category"], "safetensors_pass": r["pass_"], "reference_pass": ref_r.get("pass_"), "masked_sentence": r.get("masked_sentence", ""), } ) print("\n" + "=" * 70) print("COMPARISON AGAINST REFERENCE") print("=" * 70) print(f" Reference: {reference_path}") print(f" Total test cases: {total}") print(f" Matching pass/fail: {matching}/{total} ({pct(matching / total)})") print(f" Mismatches: {len(mismatches)}") print(f" Only in safetensors: {len(only_new)}") print(f" Only in reference: {len(only_old_ids)}") if mismatches: print("\n--- PASS/FAIL MISMATCHES ---") for m in mismatches: print( f" id={m['id']} type={m['type']} category={m['category']}" ) print( f" safetensors: {m['safetensors_pass']} reference: {m['reference_pass']}" ) print(f" sentence: {m['masked_sentence']}") equivalent = len(mismatches) == 0 print( f"\nVERDICT: {'EQUIVALENT ✓' if equivalent else 'DIFFERENCES FOUND ✗'}" ) if equivalent: print( "The safetensors (fp16) model produces identical pass/fail results\n" "to the reference fp32 checkpoint. The fp16 storage format is " "safe for Hugging Face upload." ) print() return { "total": total, "matching": matching, "mismatches": len(mismatches), "equivalent": equivalent, } def build_summary(results, model_tag="safetensors", ckpt_info=None): by_type = {} for eval_type in [ "doctrinal_association", "canonical_knowledge", "contrastive_theology", ]: matches = [r for r in results if r["type"] == eval_type] passed_matches = sum(1 for r in matches if r["pass_"]) by_type[eval_type] = { "passed": passed_matches, "total": len(matches), "pass_rate": passed_matches / len(matches) if matches else 0.0, } by_category = {} for category in sorted(set(r["category"] for r in results)): matches = [r for r in results if r["category"] == category] passed_matches = sum(1 for r in matches if r["pass_"]) by_category[category] = { "passed": passed_matches, "total": len(matches), "pass_rate": passed_matches / len(matches) if matches else 0.0, "critical_failures": sum(1 for r in matches if r["critical_failure"]), } by_difficulty = {} for difficulty in ["easy", "medium", "hard"]: matches = [r for r in results if r["difficulty"] == difficulty] passed_matches = sum(1 for r in matches if r["pass_"]) by_difficulty[difficulty] = { "passed": passed_matches, "total": len(matches), "pass_rate": passed_matches / len(matches) if matches else 0.0, } confidence_counts = dict( Counter(r["confidence"] for r in results if r["confidence"] is not None) ) return { "model_tag": model_tag, "overall_pass_rate": sum(1 for r in results if r["pass_"]) / len(results), "difficulty_weighted_score": difficulty_weighted_score(results), "critical_failure_rate": sum(1 for r in results if r["critical_failure"]) / len(results), "by_type": by_type, "by_category": by_category, "by_difficulty": by_difficulty, "confidence_counts": confidence_counts, "results": results, } # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main(): p = argparse.ArgumentParser( description="MLM evaluation for safetensors-format encoder weights" ) p.add_argument( "--repo-dir", type=str, default=".", help="Path to the HF repo directory (containing config.json, model.safetensors, and tokenizer files). Defaults to the current directory.", ) p.add_argument( "--eval-path", type=str, default="eval.json", ) p.add_argument("--device", type=str, default="cpu") p.add_argument("--k", type=int, default=5) p.add_argument( "--temperature", type=float, default=1.0, help="Logit temperature before softmax.", ) p.add_argument( "--n-samples", type=int, default=1, help="Number of multinomial draws per mask position.", ) p.add_argument("--output", type=str, default=None) p.add_argument( "--compare", type=str, default=None, help="Path to a prior eval JSON to diff pass/fail results against.", ) args = p.parse_args() device = resolve_device(args.device) repo_dir = Path(args.repo_dir).resolve() output = args.output or "eval_results/safetensors_mlm_eval.json" os.makedirs(os.path.dirname(output) or ".", exist_ok=True) print("\n" + "=" * 70) print("MLM EVALUATION SUITE — safetensors (fp16 → fp32)") print("=" * 70) print("Loading tokenizer (bert-base-uncased)...") tokenizer = BertTokenizer.from_pretrained("bert-base-uncased") print(f"Loading model from: {repo_dir}") model, config = load_model_from_repo(repo_dir, device) print( f" dim={config.n_embd}, layers={config.n_layer}, vocab={config.vocab_size}" ) # Report weight dtype status param_dtypes = set(p.dtype for p in model.parameters()) print(f" Parameter dtypes after load: {param_dtypes}") eval_path = os.path.abspath(args.eval_path) print(f"Loading eval dataset: {eval_path}...") with open(args.eval_path) as f: payload = json.load(f) test_cases = payload["test_cases"] if isinstance(payload, dict) else payload print(f"Loaded {len(test_cases)} test cases") input_tensor, attention_mask, mask_positions_list, multi_piece_flags = ( prepare_batch(test_cases, tokenizer, device) ) temp_note = ( f" (temperature={args.temperature}, n_samples={args.n_samples})" if args.temperature != 1.0 or args.n_samples > 1 else "" ) print(f"Running inference{temp_note}...") probs = run_inference( model, input_tensor, attention_mask, temperature=args.temperature ) results = score_all( test_cases, probs, mask_positions_list, multi_piece_flags, tokenizer, args.k, n_samples=args.n_samples, ) print_results(results) summary = build_summary(results) with open(output, "w") as f: json.dump(summary, f, indent=2) print(f"Saved to {output}") # Optionally compare against a reference eval if args.compare: compare_results(results, args.compare) if __name__ == "__main__": main()