{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":121150,"databundleVersionId":15833391,"isSourceIdPinned":false},{"sourceType":"datasetVersion","sourceId":15295927,"datasetId":9783371,"databundleVersionId":16199233},{"sourceType":"datasetVersion","sourceId":15022922,"datasetId":9616483,"databundleVersionId":15900503},{"sourceType":"datasetVersion","sourceId":15087303,"datasetId":9659420,"databundleVersionId":15971151},{"sourceType":"datasetVersion","sourceId":15089501,"datasetId":9652989,"databundleVersionId":15973638},{"sourceType":"datasetVersion","sourceId":14899351,"datasetId":9366631,"databundleVersionId":15763950},{"sourceType":"datasetVersion","sourceId":15249102,"datasetId":9755607,"databundleVersionId":16147110},{"sourceType":"datasetVersion","sourceId":14995108,"datasetId":9568388,"databundleVersionId":15869653},{"sourceType":"datasetVersion","sourceId":15060148,"datasetId":9641691,"databundleVersionId":15941204},{"sourceType":"datasetVersion","sourceId":15035202,"datasetId":9624709,"databundleVersionId":15913907},{"sourceType":"datasetVersion","sourceId":15099756,"datasetId":9632085,"databundleVersionId":15985064},{"sourceType":"datasetVersion","sourceId":15337452,"datasetId":9630558,"databundleVersionId":16245783},{"sourceType":"datasetVersion","sourceId":15297969,"datasetId":9785205,"databundleVersionId":16201476},{"sourceType":"datasetVersion","sourceId":15266756,"datasetId":9766214,"databundleVersionId":16166349},{"sourceType":"datasetVersion","sourceId":15060226,"datasetId":9641746,"databundleVersionId":15941290},{"sourceType":"datasetVersion","sourceId":15100435,"datasetId":9667966,"databundleVersionId":15985804},{"sourceType":"datasetVersion","sourceId":15022042,"datasetId":9615946,"databundleVersionId":15899570},{"sourceType":"datasetVersion","sourceId":15100870,"datasetId":9668265,"databundleVersionId":15986299}],"dockerImageVersionId":31286,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# !pip install evaluate sacrebleu\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n\n# ============================================\n# Deep Past Initiative – ByT5 Akkadian Probing\n# Detecting Overfitting & Memorization (v2)\n# ============================================\n#\n# Improvements over v1:\n#   - 6 probe categories (was 2): NUMBER, NAME_COPY, WORD_ORDER,\n#     SYNONYM, NOVEL_INPUT, TRAIN_VS_TEST overlap analysis\n#   - Gradient memorization score (not just binary)\n#   - Output-similarity analysis (token edit-distance to baseline)\n#   - Per-category statistics with confidence diagnostics\n#   - Test-set evaluation with train-overlap detection\n#   - Reproducible (seeded RNG)\n#   - GPU memory management between batches\n#   - Structured JSON + CSV report export\n#   - Richer novel-name pool instead of single hardcoded name\n#   - Robust early-exit on missing files\n\nimport json\nimport re\nimport random\nimport csv\nimport sys\nimport os\nfrom datetime import datetime\nfrom pathlib import Path\nfrom collections import Counter, defaultdict\nfrom dataclasses import dataclass, field, asdict\nfrom typing import List, Dict, Optional, Tuple\n\nimport torch\nimport gc\nimport pandas as pd\nfrom torch.utils.data import Dataset, DataLoader\nfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM\nfrom tqdm.auto import tqdm\n\n# =========================\n# Configuration\n# =========================\nMODEL_PATH = \"/kaggle/input/modelofinalbyt5/byt5-akkadian-model_phase1_95\"\nTRAINING_DATA_PATH = \"/kaggle/input/datasets/giovannyrodrguez/trainjson/train.jsonl\"\nTEST_DATA_PATH = \"/kaggle/input/competitions/deep-past-initiative-machine-translation/test.csv\"\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nPREFIX = \"translate Akkadian to English: \"\nMAX_SOURCE_LEN = 512\nMAX_NEW_TOKENS = 512\n\n# Generation\nBATCH_SIZE = 16\nNUM_BEAMS = 4\nEARLY_STOPPING = True\n\n# Probing\nSEED = 42\nMAX_PROBES_PER_CATEGORY = 100          # more probes → more statistical power\nMAX_TRAINING_SCAN = 5000              # how many training pairs to scan for probe material\nSIMILARITY_MEMORIZATION_THRESHOLD = 0.15  # token-edit-distance below this = \"too similar to baseline\"\n\n# Output\nREPORT_DIR = Path(\"probe_reports\")\nREPORT_DIR.mkdir(exist_ok=True)\n\nrandom.seed(SEED)\ntorch.manual_seed(SEED)\n\n# =========================\n# Novel-name pool (replaces single hardcoded name)\n# =========================\nNOVEL_NAMES = [\n    (\"pel-kur-isz-tar\", \"Pelkur-Ishtar\"),\n    (\"na-bi-dEN.ZU\", \"Nabi-Sin\"),\n    (\"tu-ra-am-dUTU\", \"Turam-Shamash\"),\n    (\"i-din-dIM\", \"Iddin-Adad\"),\n    (\"a-wi-il-dAMAR.UTU\", \"Awil-Marduk\"),\n    (\"ri-im-dEN.LÍL\", \"Rim-Enlil\"),\n    (\"pu-zur4-dINANA\", \"Puzur-Ishtar\"),\n    (\"la-ma-sa-tum\", \"Lamasatum\"),\n    (\"e-ri-ib-dUTU\", \"Erib-Shamash\"),\n    (\"ib-ni-dEN.ZU\", \"Ibni-Sin\"),\n]\n\n\n# =========================\n# Probe Utility Functions\n# =========================\n\ndef normalize_name(name: str) -> str:\n    if not name:\n        return \"\"\n    name = re.sub(r'[.,;:!?\\\"\\'\\(\\)\\[\\]]', '', name).lower()\n    replacements = {\n        'ā': 'a', 'á': 'a', 'à': 'a', 'â': 'a',\n        'ē': 'e', 'é': 'e', 'è': 'e', 'ê': 'e',\n        'ī': 'i', 'í': 'i', 'ì': 'i', 'î': 'i',\n        'ū': 'u', 'ú': 'u', 'ù': 'u', 'û': 'u',\n        'š': 'sz', 'ṣ': 's', 'ṭ': 't', 'ḫ': 'h',\n        'ʾ': '', \"'\": '',\n    }\n    for old, new in replacements.items():\n        name = name.replace(old, new)\n    return name\n\n\ndef soundex(name: str) -> str:\n    if not name:\n        return \"\"\n    name = name.upper()\n    soundex_code = name[0]\n    mapping = {\n        'B': '1', 'F': '1', 'P': '1', 'V': '1',\n        'C': '2', 'G': '2', 'J': '2', 'K': '2',\n        'Q': '2', 'S': '2', 'X': '2', 'Z': '2',\n        'D': '3', 'T': '3', 'L': '4',\n        'M': '5', 'N': '5', 'R': '6',\n    }\n    prev_code = mapping.get(name[0], '0')\n    for char in name[1:]:\n        code = mapping.get(char, '0')\n        if code != '0' and code != prev_code:\n            soundex_code += code\n        prev_code = code if code != '0' else prev_code\n        if len(soundex_code) >= 4:\n            break\n    return (soundex_code + '000')[:4]\n\n\ndef extract_potential_names(text: str) -> List[str]:\n    words = re.split(r'[\\s,;:\\(\\)\\[\\]]+', text)\n    potential = []\n    for word in words:\n        word = word.strip('.,;:!?()[]\"\\'')\n        if len(word) >= 3:\n            if word[0].isupper() or '-' in word or any(c in word for c in 'šṣṭāēīūâêîûùà'):\n                potential.append(word)\n    return potential\n\n\ndef clean_term(term: str) -> str:\n    return re.sub(r'[.,;:!?\\\"\\'\\(\\)\\[\\]]+$', '', term).strip()\n\n\ndef compute_token_edit_distance(s1: str, s2: str) -> float:\n    \"\"\"Normalized token-level edit distance ∈ [0, 1]. 0 = identical.\"\"\"\n    t1, t2 = s1.lower().split(), s2.lower().split()\n    if not t1 and not t2:\n        return 0.0\n    if not t1 or not t2:\n        return 1.0\n    m, n = len(t1), len(t2)\n    dp = [[0] * (n + 1) for _ in range(m + 1)]\n    for i in range(m + 1):\n        dp[i][0] = i\n    for j in range(n + 1):\n        dp[0][j] = j\n    for i in range(1, m + 1):\n        for j in range(1, n + 1):\n            if t1[i - 1] == t2[j - 1]:\n                dp[i][j] = dp[i - 1][j - 1]\n            else:\n                dp[i][j] = 1 + min(dp[i - 1][j], dp[i][j - 1], dp[i - 1][j - 1])\n    return dp[m][n] / max(m, n)\n\n\ndef check_term_simple(output: str, term) -> Tuple[bool, Optional[str]]:\n    if not term:\n        return False, None\n    if isinstance(term, (tuple, list)):\n        for v in term:\n            if v.lower() in output.lower():\n                return True, v\n        return False, None\n    if term.lower() in output.lower():\n        return True, term\n    return False, None\n\n\nCOMMON_WORDS_FILTER = {\n    'with', 'from', 'the', 'and', 'for', 'that', 'this', 'have', 'will',\n    'are', 'was', 'were', 'been', 'being', 'has', 'had', 'not', 'but',\n    'silver', 'gold', 'tin', 'copper', 'barley', 'oil', 'wool', 'house',\n    'son', 'daughter', 'wife', 'father', 'mother', 'brother', 'sister',\n    'seal', 'tablet', 'letter', 'witness', 'month', 'year', 'day',\n    'mina', 'shekel', 'talent', 'interest', 'debt', 'loan', 'price', 'payment',\n}\n\n\n# =========================================\n# Dataclass for structured probe results\n# =========================================\n\n@dataclass\nclass ProbeResult:\n    category: str\n    description: str\n    baseline_input: str\n    baseline_expected: str\n    baseline_output: str\n    perturbed_input: str\n    perturbed_expected: str\n    perturbed_output: str\n    baseline_check_term: str\n    perturbed_check_term: str\n    baseline_hit: bool\n    perturbed_hit: bool\n    verdict: str                     # GENERALIZED | MEMORIZED | BASELINE_FAILED\n    similarity_to_baseline: float    # token-edit-distance between perturbed_output and baseline_output\n    similarity_to_expected: float    # token-edit-distance between perturbed_output and perturbed_expected\n\n\n# =========================\n# Load Model & Tokenizer\n# =========================\n\ndef load_model(model_path: str):\n    print(f\"Loading model from {model_path} ...\")\n    try:\n        tokenizer = AutoTokenizer.from_pretrained(model_path, extra_special_tokens={})\n    except Exception as e:\n        print(f\"  ⚠ Tokenizer fallback (google/byt5-base): {e}\")\n        tokenizer = AutoTokenizer.from_pretrained(\"google/byt5-base\")\n\n    model = AutoModelForSeq2SeqLM.from_pretrained(model_path).to(DEVICE).eval()\n    print(f\"  ✓ Model on {DEVICE}  |  Parameters: {sum(p.numel() for p in model.parameters()):,}\")\n    return model, tokenizer\n\n\n# =========================\n# Inference Engine (with GPU memory management)\n# =========================\n\ndef batch_translate(\n    texts: List[str],\n    model,\n    tokenizer,\n    device,\n    batch_size: int = BATCH_SIZE,\n    show_progress: bool = True,\n) -> List[str]:\n    results = []\n    prefixed = [PREFIX + t for t in texts]\n    iterator = range(0, len(prefixed), batch_size)\n    if show_progress:\n        iterator = tqdm(iterator, desc=\"Translating\", unit=\"batch\")\n\n    for i in iterator:\n        batch = prefixed[i : i + batch_size]\n        inputs = tokenizer(\n            batch,\n            max_length=MAX_SOURCE_LEN,\n            truncation=True,\n            padding=True,\n            return_tensors=\"pt\",\n        ).to(device)\n\n        with torch.no_grad():\n            outputs = model.generate(\n                input_ids=inputs[\"input_ids\"],\n                attention_mask=inputs[\"attention_mask\"],\n                num_beams=NUM_BEAMS,\n                max_new_tokens=MAX_NEW_TOKENS,\n                early_stopping=EARLY_STOPPING,\n            )\n\n        decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)\n        results.extend([d.strip() if d.strip() else \"[empty]\" for d in decoded])\n\n        # Free GPU memory between batches\n        del inputs, outputs\n        if device.type == \"cuda\":\n            torch.cuda.empty_cache()\n\n    gc.collect()\n    return results\n\n\n# =========================\n# Load Training Data\n# =========================\n\ndef load_training_data(path: str) -> List[dict]:\n    if not os.path.exists(path):\n        print(f\"  ✗ Training file not found: {path}\")\n        sys.exit(1)\n    data = []\n    with open(path, 'r', encoding='utf-8') as f:\n        for line in f:\n            line = line.strip()\n            if line:\n                data.append(json.loads(line))\n    print(f\"  ✓ Loaded {len(data):,} training pairs\")\n    return data\n\n\ndef load_test_data(path: str) -> Optional[pd.DataFrame]:\n    if not os.path.exists(path):\n        print(f\"  ⚠ Test file not found: {path} — skipping test-set probes\")\n        return None\n    df = pd.read_csv(path)\n    print(f\"  ✓ Loaded {len(df):,} test examples\")\n    return df\n\n\n# ==========================================\n# Name Inventory Builder\n# ==========================================\n\ndef build_name_inventory(data: List[dict], max_scan: int):\n    name_inventory: Dict[str, set] = {}\n    single_name_examples = []\n    multi_name_examples = []\n\n    for p in tqdm(data[:max_scan], desc=\"Scanning names\"):\n        akk_text = str(p.get('transliteration', '')).strip()\n        eng_text = str(p.get('translation', '')).strip()\n        if len(akk_text.split()) < 4 or len(eng_text.split()) < 4:\n            continue\n\n        akk_names = extract_potential_names(akk_text)\n        eng_names = extract_potential_names(eng_text)\n        akk_sx = {soundex(n): n for n in akk_names if len(n) >= 3}\n        eng_sx = {soundex(n): n for n in eng_names if len(n) >= 3}\n        matching = set(akk_sx.keys()) & set(eng_sx.keys()) - {'A000', '0000', '', 'S416', 'S400'}\n\n        if not matching:\n            continue\n\n        matched_names = []\n        used_eng = set()\n        for sx_code in matching:\n            akk_n, eng_n = akk_sx[sx_code], eng_sx[sx_code]\n            if eng_n.lower() in COMMON_WORDS_FILTER:\n                continue\n            if eng_n[0].isupper() and eng_n not in used_eng:\n                if eng_n not in name_inventory:\n                    name_inventory[eng_n] = set()\n                name_inventory[eng_n].add(akk_n)\n                akk_m = re.search(r'\\b' + re.escape(akk_n) + r'\\b', akk_text)\n                eng_m = re.search(r'\\b' + re.escape(eng_n) + r'\\b', eng_text)\n                if akk_m and eng_m:\n                    matched_names.append({\n                        'akk_name': akk_n, 'eng_name': eng_n,\n                        'akk_pos': akk_m.start(), 'eng_pos': eng_m.start(),\n                        'akk_match': {'start': akk_m.start(), 'end': akk_m.end(), 'text': akk_n},\n                        'eng_match': {'start': eng_m.start(), 'end': eng_m.end(), 'text': eng_n},\n                    })\n                    used_eng.add(eng_n)\n\n        if len(matched_names) == 1:\n            single_name_examples.append({'akk_text': akk_text, 'eng_text': eng_text, 'name': matched_names[0]})\n        elif len(matched_names) >= 2:\n            matched_names.sort(key=lambda x: x['akk_pos'])\n            multi_name_examples.append({\n                'akk_text': akk_text, 'eng_text': eng_text,\n                'names': matched_names, 'num_names': len(matched_names),\n            })\n\n    print(f\"  ✓ Unique names: {len(name_inventory)} | Single-name sents: {len(single_name_examples)} | Multi-name: {len(multi_name_examples)}\")\n    return name_inventory, single_name_examples, multi_name_examples\n\n\n# ==========================================\n# Probe Builders\n# ==========================================\n\ndef build_number_probes(data: List[dict], max_probes: int) -> List[dict]:\n    \"\"\"Substitute a number in the source and check if the model follows.\"\"\"\n    probes = []\n    patterns = [\n        (r'(\\d+)\\s*GÍN', r'(\\d+)\\s*shekels?', 'shekel'),\n        (r'(\\d+)\\s*MA\\.NA', r'(\\d+)\\s*minas?', 'mina'),\n        (r'(\\d+)\\s*GUR', r'(\\d+)\\s*kor', 'kor'),\n        (r'(\\d+)\\s*BÁN', r'(\\d+)\\s*sūtu?', 'sutu'),\n        (r'(\\d+)\\s*sìla', r'(\\d+)\\s*liter', 'liter'),\n    ]\n    seen_sources = set()\n    for p in data[:MAX_TRAINING_SCAN]:\n        if len(probes) >= max_probes:\n            break\n        src = str(p.get('transliteration', '')).strip()\n        tgt = str(p.get('translation', '')).strip()\n        if src in seen_sources or len(src.split()) < 4:\n            continue\n        for akk_pat, eng_pat, unit in patterns:\n            am = re.search(akk_pat, src)\n            em = re.search(eng_pat, tgt)\n            if am and em:\n                onum = am.group(1)\n                if not onum.isdigit():\n                    continue\n                nnum = str(int(onum) * 2 + 1)  # odd multiplier avoids accidental collisions\n                probes.append({\n                    'category': 'NUMBER',\n                    'description': f\"NUMBER: {onum} {unit} → {nnum} {unit}\",\n                    'baseline_input': src,\n                    'baseline_expected': tgt,\n                    'baseline_check': (f\"{onum} {unit}\", onum),\n                    'perturbed_input': src[:am.start(1)] + nnum + src[am.end(1):],\n                    'perturbed_expected': tgt[:em.start(1)] + nnum + tgt[em.end(1):],\n                    'perturbed_check': (f\"{nnum} {unit}\", nnum),\n                })\n                seen_sources.add(src)\n                break\n    return probes\n\n\ndef build_name_probes(single_name_examples: List[dict], max_probes: int) -> List[dict]:\n    \"\"\"Replace a proper name with a novel (unseen) name.\"\"\"\n    probes = []\n    pool = list(NOVEL_NAMES)\n    random.shuffle(pool)\n    for i, nm_ex in enumerate(single_name_examples[:max_probes]):\n        nm = nm_ex['name']\n        src, tgt = nm_ex['akk_text'], nm_ex['eng_text']\n        novel_akk, novel_eng = pool[i % len(pool)]\n        probes.append({\n            'category': 'NAME_COPY',\n            'description': f\"NAME: {nm['eng_name']} → {novel_eng}\",\n            'baseline_input': src,\n            'baseline_expected': tgt,\n            'baseline_check': clean_term(nm['eng_name']),\n            'perturbed_input': src[:nm['akk_match']['start']] + novel_akk + src[nm['akk_match']['end']:],\n            'perturbed_expected': tgt[:nm['eng_match']['start']] + novel_eng + tgt[nm['eng_match']['end']:],\n            'perturbed_check': novel_eng,\n        })\n    return probes\n\n\ndef build_word_order_probes(data: List[dict], max_probes: int) -> List[dict]:\n    \"\"\"Swap two adjacent content words in the source. A generalising model\n    should still produce a *different* output from the unperturbed baseline\n    (or at least track the change), while a memorising model will ignore it.\"\"\"\n    probes = []\n    seen = set()\n    for p in data[:MAX_TRAINING_SCAN]:\n        if len(probes) >= max_probes:\n            break\n        src = str(p.get('transliteration', '')).strip()\n        tgt = str(p.get('translation', '')).strip()\n        tokens = src.split()\n        if len(tokens) < 6 or src in seen:\n            continue\n        # pick two adjacent content tokens (skip first/last)\n        idx = random.randint(1, len(tokens) - 3)\n        tokens[idx], tokens[idx + 1] = tokens[idx + 1], tokens[idx]\n        perturbed_src = ' '.join(tokens)\n        probes.append({\n            'category': 'WORD_ORDER',\n            'description': f\"SWAP positions {idx}↔{idx+1}\",\n            'baseline_input': src,\n            'baseline_expected': tgt,\n            'baseline_check': None,   # no specific term; we measure distance instead\n            'perturbed_input': perturbed_src,\n            'perturbed_expected': None,\n            'perturbed_check': None,\n        })\n        seen.add(src)\n    return probes\n\n\ndef build_novel_input_probes(data: List[dict], max_probes: int) -> List[dict]:\n    \"\"\"Create chimeric inputs by splicing the first half of one sentence with\n    the second half of another. A memorising model will parrot one of the two\n    training targets; a generalising model will produce something new.\"\"\"\n    probes = []\n    candidates = []\n    for p in data[:MAX_TRAINING_SCAN]:\n        src = str(p.get('transliteration', '')).strip()\n        tgt = str(p.get('translation', '')).strip()\n        if 6 <= len(src.split()) <= 30:\n            candidates.append((src, tgt))\n    random.shuffle(candidates)\n    for i in range(0, min(len(candidates) - 1, max_probes * 2), 2):\n        if len(probes) >= max_probes:\n            break\n        src_a, tgt_a = candidates[i]\n        src_b, tgt_b = candidates[i + 1]\n        mid_a = len(src_a.split()) // 2\n        mid_b = len(src_b.split()) // 2\n        chimera = ' '.join(src_a.split()[:mid_a] + src_b.split()[mid_b:])\n        probes.append({\n            'category': 'NOVEL_INPUT',\n            'description': f\"CHIMERA from pair {i} + {i+1}\",\n            'baseline_input': src_a,\n            'baseline_expected': tgt_a,\n            'baseline_check': None,\n            'perturbed_input': chimera,\n            'perturbed_expected': None,\n            'perturbed_check': None,\n            # store second parent for similarity check\n            '_parent_b_target': tgt_b,\n        })\n    return probes\n\n\ndef build_synonym_probes(data: List[dict], max_probes: int) -> List[dict]:\n    \"\"\"Replace a common Akkadian logogram with its syllabic equivalent.\n    Tests whether the model understands the *meaning* or just the surface form.\"\"\"\n    LOGOGRAM_SYNONYMS = [\n        ('KÙ.BABBAR', 'kaspum', 'silver'),\n        ('KÙ.GI', 'hurāṣum', 'gold'),\n        ('ŠE', \"še'um\", 'barley'),\n        ('GÍN', 'šiqlum', 'shekel'),\n        ('ITI', 'warahum', 'month'),\n    ]\n    probes = []\n    seen = set()\n    for p in data[:MAX_TRAINING_SCAN]:\n        if len(probes) >= max_probes:\n            break\n        src = str(p.get('transliteration', '')).strip()\n        tgt = str(p.get('translation', '')).strip()\n        if src in seen:\n            continue\n        for logo, syllabic, eng_word in LOGOGRAM_SYNONYMS:\n            if logo in src and eng_word in tgt.lower():\n                perturbed_src = src.replace(logo, syllabic, 1)\n                probes.append({\n                    'category': 'SYNONYM',\n                    'description': f\"SYNONYM: {logo} → {syllabic} (={eng_word})\",\n                    'baseline_input': src,\n                    'baseline_expected': tgt,\n                    'baseline_check': eng_word,\n                    'perturbed_input': perturbed_src,\n                    'perturbed_expected': tgt,       # same expected output\n                    'perturbed_check': eng_word,     # should still produce the english word\n                })\n                seen.add(src)\n                break\n    return probes\n\n\n# ==========================================\n# Train–Test Overlap Detection\n# ==========================================\n\ndef compute_train_test_overlap(train_data: List[dict], test_df: pd.DataFrame) -> dict:\n    \"\"\"Check for near-duplicate leakage between train and test sets.\"\"\"\n    train_sources = set()\n    for p in train_data:\n        src = str(p.get('transliteration', '')).strip()\n        if src:\n            train_sources.add(src.lower().strip())\n\n    test_col = None\n    for col in ['transliteration', 'source', 'input', 'src', 'akkadian']:\n        if col in test_df.columns:\n            test_col = col\n            break\n    if test_col is None:\n        # try first column\n        test_col = test_df.columns[0]\n\n    exact_matches = 0\n    high_overlap = 0\n    total = 0\n    for _, row in test_df.iterrows():\n        test_src = str(row[test_col]).strip()\n        if not test_src:\n            continue\n        total += 1\n        if test_src.lower().strip() in train_sources:\n            exact_matches += 1\n        else:\n            # token-level Jaccard overlap\n            test_tokens = set(test_src.lower().split())\n            best_jaccard = 0.0\n            for ts in random.sample(list(train_sources), min(2000, len(train_sources))):\n                train_tokens = set(ts.split())\n                inter = len(test_tokens & train_tokens)\n                union = len(test_tokens | train_tokens)\n                if union > 0:\n                    j = inter / union\n                    if j > best_jaccard:\n                        best_jaccard = j\n            if best_jaccard > 0.8:\n                high_overlap += 1\n\n    return {\n        'total_test': total,\n        'exact_matches': exact_matches,\n        'high_overlap_samples': high_overlap,\n        'exact_match_pct': round(exact_matches / max(total, 1) * 100, 2),\n        'high_overlap_pct': round(high_overlap / max(total, 1) * 100, 2),\n    }\n\n\n# ==========================================\n# Main Execution\n# ==========================================\n\ndef main():\n    timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n\n    # ---- Load data ----\n    print(\"\\n📂 Loading data...\")\n    train_data = load_training_data(TRAINING_DATA_PATH)\n    test_df = load_test_data(TEST_DATA_PATH)\n\n    # ---- Load model ----\n    print(\"\\n🤖 Loading model...\")\n    model, tokenizer = load_model(MODEL_PATH)\n\n    # ---- Build name inventory ----\n    print(\"\\n🔎 Building name inventory...\")\n    name_inventory, single_name_ex, multi_name_ex = build_name_inventory(train_data, MAX_TRAINING_SCAN)\n\n    # ---- Build probes ----\n    print(\"\\n🏗️  Building probes...\")\n    all_probes = []\n    all_probes += build_number_probes(train_data, MAX_PROBES_PER_CATEGORY)\n    all_probes += build_name_probes(single_name_ex, MAX_PROBES_PER_CATEGORY)\n    all_probes += build_word_order_probes(train_data, MAX_PROBES_PER_CATEGORY)\n    all_probes += build_novel_input_probes(train_data, MAX_PROBES_PER_CATEGORY)\n    all_probes += build_synonym_probes(train_data, MAX_PROBES_PER_CATEGORY)\n\n    cat_counts = Counter(p['category'] for p in all_probes)\n    print(f\"  ✓ Total probes: {len(all_probes)}\")\n    for cat, cnt in sorted(cat_counts.items()):\n        print(f\"      {cat}: {cnt}\")\n\n    # ---- Run inference ----\n    print(f\"\\n🚀 Running {len(all_probes)} probes (×2 = {len(all_probes)*2} translations)...\")\n    inputs = []\n    for p in all_probes:\n        inputs.append(p['baseline_input'])\n        inputs.append(p['perturbed_input'])\n\n    outputs = batch_translate(inputs, model, tokenizer, DEVICE)\n\n    # ---- Evaluate ----\n    print(\"\\n📊 Evaluating results...\")\n    results: List[ProbeResult] = []\n    category_stats = defaultdict(lambda: {'gen': 0, 'mem': 0, 'fail': 0, 'sim_scores': []})\n\n    for i, p in enumerate(all_probes):\n        bo = outputs[i * 2]\n        po = outputs[i * 2 + 1]\n\n        # Similarity metrics\n        sim_to_baseline = 1.0 - compute_token_edit_distance(po, bo)\n        sim_to_expected = 0.0\n        if p.get('perturbed_expected'):\n            sim_to_expected = 1.0 - compute_token_edit_distance(po, p['perturbed_expected'])\n\n        # Term checks\n        if p['baseline_check'] is not None:\n            bp, _ = check_term_simple(bo, p['baseline_check'])\n        else:\n            bp = True  # no term to check → baseline always \"passes\"\n\n        if p['perturbed_check'] is not None:\n            pp, _ = check_term_simple(po, p['perturbed_check'])\n        else:\n            # For probes without a check-term (WORD_ORDER, NOVEL_INPUT),\n            # a \"generalised\" model should produce output *different* from baseline\n            pp = sim_to_baseline < (1.0 - SIMILARITY_MEMORIZATION_THRESHOLD)\n\n        if bp and pp:\n            verdict = \"GENERALIZED\"\n            category_stats[p['category']]['gen'] += 1\n        elif bp and not pp:\n            verdict = \"MEMORIZED\"\n            category_stats[p['category']]['mem'] += 1\n        else:\n            verdict = \"BASELINE_FAILED\"\n            category_stats[p['category']]['fail'] += 1\n\n        category_stats[p['category']]['sim_scores'].append(sim_to_baseline)\n\n        results.append(ProbeResult(\n            category=p['category'],\n            description=p['description'],\n            baseline_input=p['baseline_input'][:120],\n            baseline_expected=(p['baseline_expected'] or '')[:120],\n            baseline_output=bo[:120],\n            perturbed_input=p['perturbed_input'][:120],\n            perturbed_expected=(p.get('perturbed_expected') or '')[:120],\n            perturbed_output=po[:120],\n            baseline_check_term=str(p['baseline_check'] or ''),\n            perturbed_check_term=str(p['perturbed_check'] or ''),\n            baseline_hit=bp,\n            perturbed_hit=pp,\n            verdict=verdict,\n            similarity_to_baseline=round(sim_to_baseline, 4),\n            similarity_to_expected=round(sim_to_expected, 4),\n        ))\n\n    # ---- Train/Test overlap (if test data available) ----\n    overlap_info = None\n    if test_df is not None:\n        print(\"\\n🔍 Checking train–test overlap...\")\n        overlap_info = compute_train_test_overlap(train_data, test_df)\n\n    # ---- Console Report ----\n    print(\"\\n\" + \"=\" * 60)\n    print(\"                    PROBE REPORT\")\n    print(\"=\" * 60)\n\n    total_gen = sum(s['gen'] for s in category_stats.values())\n    total_mem = sum(s['mem'] for s in category_stats.values())\n    total_fail = sum(s['fail'] for s in category_stats.values())\n    total_valid = total_gen + total_mem\n\n    for cat in sorted(category_stats.keys()):\n        s = category_stats[cat]\n        n = s['gen'] + s['mem'] + s['fail']\n        valid = s['gen'] + s['mem']\n        gen_rate = s['gen'] / valid * 100 if valid > 0 else 0\n        avg_sim = sum(s['sim_scores']) / len(s['sim_scores']) if s['sim_scores'] else 0\n        print(f\"\\n  [{cat}]  (n={n})\")\n        print(f\"    Generalized : {s['gen']:3d}  |  Memorized : {s['mem']:3d}  |  Baseline fail : {s['fail']:3d}\")\n        if valid > 0:\n            print(f\"    Gen rate    : {gen_rate:5.1f}%\")\n        print(f\"    Avg similarity to baseline (perturbed output): {avg_sim:.3f}\")\n\n    print(f\"\\n  {'─'*50}\")\n    print(f\"  OVERALL  (valid probes: {total_valid})\")\n    if total_valid > 0:\n        print(f\"    Generalization rate : {total_gen/total_valid*100:.1f}%\")\n        print(f\"    Memorization rate   : {total_mem/total_valid*100:.1f}%\")\n    print(f\"    Baseline failures   : {total_fail}\")\n\n    if overlap_info:\n        print(f\"\\n  TRAIN–TEST OVERLAP\")\n        print(f\"    Test examples     : {overlap_info['total_test']}\")\n        print(f\"    Exact duplicates  : {overlap_info['exact_matches']} ({overlap_info['exact_match_pct']}%)\")\n        print(f\"    High overlap (>80%): {overlap_info['high_overlap_samples']} ({overlap_info['high_overlap_pct']}%)\")\n\n    print(\"=\" * 60)\n\n    # ---- Export ----\n    # JSON report\n    report = {\n        'timestamp': timestamp,\n        'model_path': MODEL_PATH,\n        'device': str(DEVICE),\n        'num_probes': len(all_probes),\n        'category_stats': {\n            cat: {\n                'generalized': s['gen'],\n                'memorized': s['mem'],\n                'baseline_failed': s['fail'],\n                'gen_rate': round(s['gen'] / max(s['gen'] + s['mem'], 1) * 100, 2),\n                'avg_sim_to_baseline': round(sum(s['sim_scores']) / max(len(s['sim_scores']), 1), 4),\n            }\n            for cat, s in category_stats.items()\n        },\n        'overall': {\n            'generalized': total_gen,\n            'memorized': total_mem,\n            'baseline_failed': total_fail,\n            'gen_rate': round(total_gen / max(total_valid, 1) * 100, 2),\n        },\n        'train_test_overlap': overlap_info,\n    }\n    json_path = REPORT_DIR / f\"probe_report_{timestamp}.json\"\n    with open(json_path, 'w') as f:\n        json.dump(report, f, indent=2)\n    print(f\"\\n💾 JSON report → {json_path}\")\n\n    # CSV detail\n    csv_path = REPORT_DIR / f\"probe_details_{timestamp}.csv\"\n    with open(csv_path, 'w', newline='', encoding='utf-8') as f:\n        writer = csv.DictWriter(f, fieldnames=[\n            'category', 'description', 'verdict',\n            'similarity_to_baseline', 'similarity_to_expected',\n            'baseline_hit', 'perturbed_hit',\n            'baseline_output', 'perturbed_output',\n            'perturbed_input',\n        ])\n        writer.writeheader()\n        for r in results:\n            writer.writerow({\n                'category': r.category,\n                'description': r.description,\n                'verdict': r.verdict,\n                'similarity_to_baseline': r.similarity_to_baseline,\n                'similarity_to_expected': r.similarity_to_expected,\n                'baseline_hit': r.baseline_hit,\n                'perturbed_hit': r.perturbed_hit,\n                'baseline_output': r.baseline_output,\n                'perturbed_output': r.perturbed_output,\n                'perturbed_input': r.perturbed_input,\n            })\n    print(f\"💾 CSV details → {csv_path}\")\n\n    # ---- Diagnostic flags ----\n    print(\"\\n🩺 Diagnostic flags:\")\n    if total_valid > 0 and total_mem / total_valid > 0.5:\n        print(\"  🔴 HIGH MEMORIZATION — model memorises >50% of probes\")\n    elif total_valid > 0 and total_mem / total_valid > 0.25:\n        print(\"  🟡 MODERATE MEMORIZATION — model memorises 25–50% of probes\")\n    else:\n        print(\"  🟢 LOW MEMORIZATION — model generalises well on probed categories\")\n\n    if total_fail > total_valid * 0.3:\n        print(\"  🟡 HIGH BASELINE FAILURE — model struggles with baseline inputs; probes may be unreliable\")\n\n    if overlap_info and overlap_info['exact_match_pct'] > 5:\n        print(\"  🔴 DATA LEAKAGE — significant train/test overlap detected\")\n\n    print(\"\\nDone.\")\n\n\nif __name__ == \"__main__\":\n    main()\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}