set -a
[ -f .env ] && . ./.env
set +a

if [ -z "${NGC_API_KEY:-}" ] && [ -n "${NVIDIA_API_KEY:-}" ]; then
  export NGC_API_KEY="$NVIDIA_API_KEY"
fi
: "${NGC_API_KEY:?Set NGC_API_KEY or NVIDIA_API_KEY}"
: "${LOCAL_NIM_CACHE:?Set LOCAL_NIM_CACHE}"

echo "$NGC_API_KEY" | docker login nvcr.io --username '$oauthtoken' --password-stdin

# 40B default: 0,1 for 2x H100; set 0 for a single H200.
export NIM_TEST_GPUS="${NIM_TEST_GPUS:-0,1}"
mkdir -p "${LOCAL_NIM_CACHE}"
chmod 700 "${LOCAL_NIM_CACHE}"   # owner-only; if the NIM runs as a different UID, add -u "$(id -u)" to docker run

# For 7B: export NIM_VARIANT=7b; export NIM_TEST_GPUS="${NIM_TEST_GPUS:-0}"
docker run --rm -it --name evo2-nim \
  --runtime=nvidia \
  --gpus "\"device=${NIM_TEST_GPUS}\"" \
  -e NGC_API_KEY \
  -e NIM_VARIANT \
  -v "${LOCAL_NIM_CACHE}:/opt/nim/.cache" \
  -p 8000:8000 \
  nvcr.io/nim/arc/evo2:2

until curl -sf http://localhost:8000/v1/health/ready; do sleep 10; done

import json
import os
from pathlib import Path
import requests

HOSTED = True

def clean_dna(value: str) -> str:
    seq = "".join(value.upper().split())
    invalid = sorted(set(seq) - set("ACGT"))
    if invalid:
        raise ValueError(f"Unexpected DNA characters: {''.join(invalid)}")
    return seq

prompt = clean_dna("ACTGACTGACTGACTG")
nim_url = os.getenv("EVO2_NIM_URL", "http://localhost:8000")
url = (
    "https://health.api.nvidia.com/v1/biology/arc/evo2-40b/generate"
    if HOSTED else f"{nim_url}/biology/arc/evo2/generate"
)
headers = {"Content-Type": "application/json"}
if HOSTED:
    api_key = os.getenv("NGC_API_KEY")
    headers["Authorization"] = f"Bearer {api_key}"

payload = {
    "sequence": prompt,
    "num_tokens": 64,
    "temperature": 0.7,
    "top_k": 3,
    "top_p": 0.0,
    "random_seed": 1,
    "enable_sampled_probs": True,
    "enable_elapsed_ms_per_token": True,
}
response = requests.post(url, headers=headers, json=payload, timeout=180)
response.raise_for_status()
result = response.json()
seq = result["sequence"]
if sorted(set(seq.upper()) - set("ACGT")):
    raise ValueError("Generated sequence contains unexpected non-ACGT bases")

Path("evo2_generation.json").write_text(json.dumps(result, indent=2) + "\n")
Path("evo2_generated.fa").write_text(f">evo2_generated\n{seq}\n")
print(f"Generated {len(seq)} bases in {result.get('elapsed_ms')} ms")

import base64
import io
import os
import numpy as np
import requests

nim_url = os.getenv("EVO2_NIM_URL", "http://localhost:8000")
payload = {
    "sequence": clean_dna("ACTGACTGACTG"),
    "output_layers": ["output_layer", "decoder.layers.3.self_attention"],
}
response = requests.post(
    f"{nim_url}/biology/arc/evo2/forward",
    headers={"Content-Type": "application/json"},
    json=payload,
    timeout=300,
)
response.raise_for_status()
npz_bytes = base64.b64decode(response.json()["data"])
with open("evo2_forward_outputs.npz", "wb") as handle:
    handle.write(npz_bytes)
arrays = np.load(io.BytesIO(npz_bytes), allow_pickle=False)
for name in arrays.files:
    arr = arrays[name]
    print(name, arr.shape, arr.dtype, bool(np.isfinite(arr).all()), float(arr.mean()))