=== BIP_FAST_FINAL.ipynb ===

# ---
#@title 1. Setup and Install Dependencies { display-mode: "form" }
#@markdown Installs packages and detects GPU/TPU. Memory-optimized for Colab.

print("=" * 60)
print("BIP TEMPORAL INVARIANCE EXPERIMENT")
print("=" * 60)
print()

# Progress tracker
TASKS = [
    "Install dependencies",
    "Clone Sefaria corpus (~8GB)",
    "Clone sqnd-probe repo (Dear Abby data)",
    "Preprocess corpora",
    "Extract bond structures",
    "Generate train/test splits",
    "Train BIP model",
    "Evaluate res

# ---
#@title 4. Define Data Classes and Loaders { display-mode: "form" }
#@markdown Defines enums, dataclasses, and corpus loaders.

import json
import hashlib
import re
from dataclasses import dataclass, field, asdict
from typing import List, Dict
from enum import Enum
from collections import defaultdict
from tqdm.auto import tqdm

print("Defining data structures...")

class TimePeriod(Enum):
    BIBLICAL = 0        # ~1000-500 BCE
    SECOND_TEMPLE = 1   # ~500 BCE - 70 CE
    TANNAITIC = 2       #

# ---
#@title 6. Extract Bond Structures { display-mode: "form" }
#@markdown Extracts moral bond structures. Streams to disk to save memory.

import gc

mark_task("Extract bond structures", "running")

RELATION_PATTERNS = {
    BondType.HARM_PREVENTION: [r'\b(kill|murder|harm|hurt|save|rescue|protect|danger)\b'],
    BondType.RECIPROCITY: [r'\b(return|repay|owe|debt|mutual|exchange)\b'],
    BondType.AUTONOMY: [r'\b(choose|decision|consent|agree|force|coerce|right)\b'],
    BondType.PROPERTY: [r'\b(pr

# ---
#@title 8. Define BIP Model Architecture { display-mode: "form" }
#@markdown Defines the model. Clears memory first to avoid OOM.

import gc
import psutil

# CRITICAL: Clear memory before loading model
print("Clearing memory before model load...")
gc.collect()

if torch.cuda.is_available():
    torch.cuda.empty_cache()

mem = psutil.virtual_memory()
print(f"Memory before model: {mem.used/1e9:.1f}/{mem.total/1e9:.1f} GB ({mem.percent}%)")
print()


import torch
import torch.nn as nn
import torch.

# ---
#@title 9. Train BIP Model - BIDIRECTIONAL { display-mode: "form" }
#@markdown Trains on BOTH directions for stronger invariance testing.

import gc
import psutil

mark_task("Train BIP model", "running")

# Memory cleanup before training
gc.collect()
if torch.cuda.is_available():
    torch.cuda.empty_cache()

mem = psutil.virtual_memory()
print(f"Memory at start: {mem.used/1e9:.1f}/{mem.total/1e9:.1f} GB")

print("=" * 60)
print("BIDIRECTIONAL BIP TRAINING")
print("=" * 60)
print()
print(f"Accel

# ---
#@title 10. Evaluate Bidirectional Results { display-mode: "form" }
#@markdown Compares results from BOTH directions to assess true invariance.

import gc
import psutil

mark_task("Evaluate results", "running")

mem = psutil.virtual_memory()
print(f"Memory at eval start: {mem.used/1e9:.1f}/{mem.total/1e9:.1f} GB")

print("=" * 60)
print("BIDIRECTIONAL BIP RESULTS")
print("=" * 60)
print()

chance_time = 1/9  # 9 time periods
chance_hohfeld = 1/4  # 4 Hohfeld states

print("DIRECTION A: Ancient →

# ---
#@title 11. Download Results (Optional) { display-mode: "form" }
#@markdown Creates a zip file with model checkpoint and metrics.

import shutil
from google.colab import files

# Create results directory
!mkdir -p results
!cp models/checkpoints/best_model.pt results/
!cp data/splits/all_splits.json results/

# Save metrics
if len(train_dataset) > 0:
    metrics = {
        'accelerator': ACCELERATOR,
        'time_acc_from_bond': time_acc,
        'hohfeld_acc': hohfeld_acc,
        'chance_leve

