social-data-attribution-w2/scripts/enrichment/debug_enrich.py

"""Diagnostic script: run on the cluster to debug the enrichment CUDA crash.

Usage (on the cluster, inside an interactive GPU session or via sbatch):
    CUDA_LAUNCH_BLOCKING=1 HF_DATASETS_TRUST_REMOTE_CODE=1 \
        uv run python scripts/enrichment/debug_enrich.py
"""

import json
import sys
import torch


def main():
    print(f"Python: {sys.version}")
    print(f"PyTorch: {torch.__version__}")
    print(f"CUDA available: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"CUDA device: {torch.cuda.get_device_name(0)}")

    import transformers

    print(f"Transformers: {transformers.__version__}")

    # --- Step 1: Load tokenizer and model on CPU first ---
    print("\n=== Step 1: Loading tokenizer & model on CPU ===")
    from transformers import AutoTokenizer, AutoModelForSequenceClassification

    model_name = "WebOrganizer/FormatClassifier-NoURL"

    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    print(f"Tokenizer vocab size: {tokenizer.vocab_size}")
    print(f"Tokenizer type: {type(tokenizer).__name__}")

    model = AutoModelForSequenceClassification.from_pretrained(
        model_name,
        trust_remote_code=True,
        use_memory_efficient_attention=False,
        unpad_inputs=False,
        torch_dtype=torch.bfloat16,
    )
    print(f"Model config vocab_size: {model.config.vocab_size}")
    print(f"Model config unpad_inputs: {model.config.unpad_inputs}")
    print(
        f"Model config use_memory_efficient_attention: {model.config.use_memory_efficient_attention}"
    )
    print(
        f"Model embedding weight shape: {model.new.embeddings.word_embeddings.weight.shape}"
    )

    # --- Step 2: Test tokenization ---
    print("\n=== Step 2: Tokenize sample texts ===")
    data_path = "data/dolma3_mix_first.jsonl"
    texts = []
    try:
        with open(data_path) as f:
            for i, line in enumerate(f):
                if i >= 20:
                    break
                record = json.loads(line)
                text = record.get("text", "")
                texts.append(text[:2000])  # truncate for speed
    except FileNotFoundError:
        print(f"WARNING: {data_path} not found, using synthetic test texts")
        texts = [
            "Hello world, this is a test.",
            "https://example.com\n\nSome content here with special chars: <>&",
            "A" * 5000,  # very long text
            "",  # empty text
            "\x00\x01\x02",  # control characters
        ]

    batch = tokenizer(
        texts,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=1024,
    )
    input_ids = batch["input_ids"]
    print(f"Batch input_ids shape: {input_ids.shape}")
    print(f"Min token ID: {input_ids.min().item()}")
    print(f"Max token ID: {input_ids.max().item()}")
    print(f"Model vocab size: {model.config.vocab_size}")

    if input_ids.max().item() >= model.config.vocab_size:
        print(
            f"ERROR: Token ID {input_ids.max().item()} >= vocab_size {model.config.vocab_size}!"
        )
        # Find which text caused it
        for i in range(input_ids.shape[0]):
            row_max = input_ids[i].max().item()
            if row_max >= model.config.vocab_size:
                print(f"  Row {i} has max token ID {row_max}")
                print(f"  Text preview: {texts[i][:200]!r}")
        return

    print("Token IDs are within vocab range ✓")

    # --- Step 3: Run inference on CPU ---
    print("\n=== Step 3: CPU inference ===")
    model.eval()
    try:
        with torch.inference_mode():
            outputs = model(**batch)
        print(f"CPU inference OK! Output logits shape: {outputs.logits.shape}")
    except Exception as e:
        print(f"CPU inference FAILED: {type(e).__name__}: {e}")
        import traceback

        traceback.print_exc()
        return

    # --- Step 4: Run inference on GPU ---
    if not torch.cuda.is_available():
        print("\nNo GPU available, skipping GPU test.")
        return

    print("\n=== Step 4: GPU inference (CUDA_LAUNCH_BLOCKING=1 recommended) ===")
    model = model.to("cuda")
    batch_gpu = {k: v.to("cuda") for k, v in batch.items()}
    try:
        with torch.inference_mode():
            outputs = model(**batch_gpu)
        print(f"GPU inference OK! Output logits shape: {outputs.logits.shape}")
    except Exception as e:
        print(f"GPU inference FAILED: {type(e).__name__}: {e}")
        import traceback

        traceback.print_exc()
        return

    # --- Step 5: Test with the full FormatClassifier wrapper ---
    print("\n=== Step 5: Full FormatClassifier test ===")
    from dolma.format_model import FormatClassifier

    classifier = FormatClassifier(
        device="cuda",
        max_length=1024,
        torch_dtype=torch.bfloat16,
        use_memory_efficient_attention=False,
        unpad_inputs=False,
        use_nourl_fallback=True,
    )
    print(
        f"FormatClassifier loaded. Config unpad_inputs: {classifier.model.config.unpad_inputs}"
    )

    test_urls = [None, "https://example.com", None]
    test_texts = texts[:3] if len(texts) >= 3 else texts + ["test"] * (3 - len(texts))
    try:
        probs, labels = classifier.predict_batch(
            test_urls[: len(test_texts)], test_texts
        )
        print(f"predict_batch OK! Labels: {labels}")
    except Exception as e:
        print(f"predict_batch FAILED: {type(e).__name__}: {e}")
        import traceback

        traceback.print_exc()

    print("\n=== All diagnostics complete ===")


if __name__ == "__main__":
    main()