CardioLLM

This guide explains how to download the ECG-CoCa encoder and load the CardioLLM checkpoint from Hugging Face.

1. Download the ECG-CoCa checkpoint

Download the ECG-CoCa checkpoint from Google Drive

2. Load CardioLLM from Hugging Face

from pathlib import Path

import torch
from transformers import AutoTokenizer
from llava.model.language_model.llava_llama import (
    LlavaConfig,
    LlavaLlamaForCausalLM,
)

MODEL_ID = "klyang/CardioLLM"  # Replace with your repository ID.
REVISION = "main"  # Use a full commit SHA to pin a specific model version.
ECG_CHECKPOINT = Path("checkpoints/cpt_wfep_epoch_20.pt").resolve()


device = "cuda:0" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device.startswith("cuda") else torch.float32
print(f"Loading {MODEL_ID} on {device} with {dtype}", flush=True)

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    revision=REVISION,
    use_fast=False,
)
config = LlavaConfig.from_pretrained(MODEL_ID, revision=REVISION)
if config.model_type != "llava_llama":
    raise ValueError(f"Unexpected architecture: {config.model_type}")

# Override the local encoder path without changing the published config.json.
config.mm_ecg_tower = str(ECG_CHECKPOINT)
# Instantiate the vision encoder before restoring its fine-tuned weights.
config.unfreeze_mm_vision_tower = True

model, loading_info = LlavaLlamaForCausalLM.from_pretrained(
    MODEL_ID,
    revision=REVISION,
    config=config,
    torch_dtype=dtype,
    device_map={"": device},
    low_cpu_mem_usage=True,
    use_safetensors=True,
    attn_implementation="eager",
    output_loading_info=True,
)

# Do not silently accept missing or randomly initialized parameters.
for category in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs"):
    problems = loading_info.get(category, [])
    if problems:
        raise RuntimeError(f"Incomplete checkpoint load: {category}: {problems[:10]}")

unloaded = [name for name, parameter in model.named_parameters() if parameter.is_meta]
if unloaded:
    raise RuntimeError(f"Unloaded parameters: {unloaded[:10]}")

for name, tower in (
    ("ECG", model.get_ecg_tower()),
    ("vision", model.get_vision_tower()),
):
    loaded = getattr(tower, "is_loaded", False)
    loaded = loaded() if callable(loaded) else loaded
    if tower is None or not loaded:
        raise RuntimeError(f"The {name} encoder was not initialized.")

model.eval()
print("CardioLLM, tokenizer, ECG encoder, and vision encoder loaded successfully.")

# A small text-only forward pass checks that the loaded language model runs.
batch = tokenizer("Hello.", return_tensors="pt")
inputs = {
    key: value.to(device)
    for key, value in batch.items()
    if key in {"input_ids", "attention_mask"}
}
with torch.inference_mode():
    logits = model(**inputs, use_cache=False, return_dict=True).logits

expected_shape = (*inputs["input_ids"].shape, config.vocab_size)
if tuple(logits.shape) != expected_shape or not torch.isfinite(logits).all().item():
    raise RuntimeError("The forward pass returned invalid logits.")
print(f"Forward pass passed: logits shape = {tuple(logits.shape)}")