probe_evo2.py

#!/usr/bin/env python
"""Probe: can we load Evo 2 in a job container and get (a) per-token logits/NLL
and (b) last-layer embeddings? Prints everything; try/except everywhere."""
import sys
print("python", sys.version, flush=True)
import torch
try:
    from transformer_engine.common.recipe import _OverrideLinearPrecision
    torch.serialization.add_safe_globals([_OverrideLinearPrecision])
    print("TE safe_globals registered", flush=True)
except Exception as e:
    print("TE safe_globals skipped:", str(e)[:120], flush=True)
print("torch", torch.__version__, "cuda:", torch.cuda.is_available(),
      torch.cuda.get_device_name(0) if torch.cuda.is_available() else "", flush=True)

from evo2 import Evo2
print("evo2 package imported", flush=True)
m = Evo2("evo2_1b_base")
print("model loaded: evo2_1b_base", flush=True)
print("attrs:", [x for x in dir(m) if not x.startswith("_")], flush=True)

seq = "ACGTGCTAGCTAGCATCGATCGTGCATGCATCGATCGTACGATCG" * 8

# --- (a) scoring ---
for call in [
    lambda: m.score_sequences([seq]),
    lambda: m.score_sequences([seq], batch_size=1),
]:
    try:
        print("score_sequences ->", call(), flush=True)
        break
    except Exception as e:
        print("score attempt failed:", type(e).__name__, str(e)[:200], flush=True)

# --- (b) tokenization + forward for logits/hidden states ---
try:
    toks = m.tokenizer.tokenize(seq)
    print("n tokens:", len(toks), flush=True)
except Exception as e:
    print("tokenizer failed:", str(e)[:200], flush=True)
    toks = None

if toks is not None:
    try:
        ids = torch.tensor([toks], dtype=torch.long).cuda()
        core = m.model
        print("core type:", type(core).__name__, flush=True)
        with torch.no_grad():
            out = core(ids)
        print("forward out type:", type(out), flush=True)
        if isinstance(out, tuple):
            for i, o in enumerate(out):
                try: print(f"out[{i}] shape:", tuple(o.shape), flush=True)
                except Exception: print(f"out[{i}] type:", type(o), flush=True)
        else:
            try: print("out shape:", tuple(out.shape), flush=True)
            except Exception: print("out:", str(out)[:300], flush=True)
    except Exception as e:
        print("forward failed:", type(e).__name__, str(e)[:300], flush=True)
print("PROBE DONE", flush=True)