#!/usr/bin/env python3
"""Probe the HF DeepSeekV4 indexer implementation to understand the correct architecture.
Specifically: what shape are the indexer compressed keys, and how does scoring work?
Run via: fire_b200_test probe_hf_indexer.py
"""
import sys, os

# Find the HF modeling file
candidates = [
    "/root/dsv4-nvfp4-workspace/venv/lib/python3.12/site-packages/transformers/models/deepseek_v4/modeling_deepseek_v4.py",
    "/root/dsv4-nvfp4-workspace/venv/lib/python*/site-packages/transformers/models/deepseek_v4/modeling_deepseek_v4.py",
]

# Also try to find it dynamically
import glob
matches = glob.glob("/root/dsv4-nvfp4-workspace/venv/lib/python*/site-packages/transformers/models/deepseek_v4/modeling_deepseek_v4.py")
if matches:
    candidates = matches

found = None
for c in candidates:
    if os.path.exists(c):
        found = c
        break

if found is None:
    # Try pip show
    import subprocess
    result = subprocess.run(["find", "/root/dsv4-nvfp4-workspace/venv", "-name", "modeling_deepseek_v4.py"], 
                          capture_output=True, text=True)
    if result.stdout.strip():
        found = result.stdout.strip().split('\n')[0]

if found:
    print(f"Found: {found}")
    # Read and print the indexer-related code
    with open(found) as f:
        lines = f.readlines()
    
    # Find class definitions and indexer-related methods
    in_relevant = False
    indent = 0
    for i, line in enumerate(lines):
        # Look for indexer, compress, lightning, score keywords
        lower = line.lower()
        if any(kw in lower for kw in ['indexer', 'lightning', 'index_score', 'index_topk', 'compress_indexer', 'indexer_head']):
            # Print surrounding context
            start = max(0, i - 2)
            end = min(len(lines), i + 20)
            print(f"\n--- Line {i+1} ---")
            for j in range(start, end):
                marker = ">>>" if j == i else "   "
                print(f"{marker} {j+1}: {lines[j]}", end='')
else:
    print("DeepSeek V4 modeling file not found. Checking what's available...")
    result = subprocess.run(["find", "/root/dsv4-nvfp4-workspace/venv", "-name", "modeling_deepseek*.py"], 
                          capture_output=True, text=True)
    print(result.stdout[:2000] if result.stdout else "No deepseek modeling files found")
    
    # Try pip
    result2 = subprocess.run(["pip", "show", "transformers"], capture_output=True, text=True)
    print(result2.stdout[:500])

print("\nDone.")