Deploy with bundled camel_inference and LFS for demo data

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.dat filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,4 +1,3 @@
-camel_inference/
 checkpoints/
 __pycache__/
 *.pyc

camel_inference/.gitignore

@@ -0,0 +1,12 @@
+# MacOS
+.DS_Store
+
+# Python
+/.env
+__pycache__
+
+# Ignore model checkpoints
+checkpoints/
+*.pt
+
+*.egg-info

camel_inference/README.md

@@ -0,0 +1,82 @@
+# CAMEL Inference
+
+Inference-only repository for running CAMEL ECG-language checkpoints.
+
+Only `run_camel.py` is intended as a public entrypoint. Modules under `src/camel/` are internal implementation details and may change.
+
+## Repository Layout
+
+- `run_camel.py`: public inference CLI
+- `src/camel/`: internal model, tokenizer, ECG packing, and loading utilities
+- `checkpoints/`: local adapter/checkpoint files
+
+## Requirements
+
+- Python 3.10+
+- CUDA-enabled PyTorch recommended for practical inference latency
+
+## Install
+
+```bash
+conda create -n camel python=3.10 -y
+conda activate camel
+pip install -e .
+```
+
+## Checkpoints
+
+Checkpoints must be downloaded from huggingface `CAMEL-ECG/CAMEL` or with the repository script:
+
+```bash
+bash scripts/download_checkpoints.sh
+```
+
+## Usage
+
+* CAMEL is available in three modes: 
+  - `base`
+  - `ecgbench`
+  - `forecast`
+
+  ```bash
+  python run_camel.py \
+    --mode forecast \
+    --text "Forecast cardiac rhythm for the next 5 minutes." \
+    --ecgs demo/08704_hr \
+    --device cuda:0
+  ```
+
+  ```bash
+  python run_camel.py \
+    --mode base \
+    --text "Compare the two ECG waveforms." \
+    --ecgs demo/12585_hr demo/12646_hr \
+    --device cuda:0
+  ```
+
+* Optionally, you can set start, end, and leads with `--ecgs-config`.
+
+  ```bash
+  python run_camel.py \
+    --mode forecast \
+    --text "Forecast cardiac rhythm for the next 5 minutes." \
+    --ecgs demo/08704_hr \
+    --ecg-configs "start:0;end:5;use_leads:I,II" \
+    --device cuda:0
+  ```
+
+* Using `--text` and `--ecgs` defaults to text followed by the ecg in order. 
+For arbitrary text/ECG interleaving use `--json`.
+  ```bash
+  python run_camel.py --mode base --json demo/example_prompt.json --device cuda:0
+  ```
+
+* Sampling flags:
+  - `--temperature`
+  - `--top-k`
+  - `--top-p`
+  - `--min-p`
+  - `--max-new-tokens`
+
+Implementation notes:
+- ECG loading is currently implemented for WFDB-format inputs. To support additional formats, extend `src/read_ecg.py`.

camel_inference/demo/08704_hr.dat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ade700b2d1a1e3a0889d4135ee179a33b7dc2eaf76d6702bc9bc64f6cef4a3f4
+size 120000

camel_inference/demo/08704_hr.hea

@@ -0,0 +1,13 @@
+08704_hr 12 500 5000
+08704_hr.dat 16 1000.0(0)/mV 16 0 -50 41247 0 I
+08704_hr.dat 16 1000.0(0)/mV 16 0 -20 64030 0 II
+08704_hr.dat 16 1000.0(0)/mV 16 0 30 22786 0 III
+08704_hr.dat 16 1000.0(0)/mV 16 0 35 13135 0 aVR
+08704_hr.dat 16 1000.0(0)/mV 16 0 -40 42425 0 aVL
+08704_hr.dat 16 1000.0(0)/mV 16 0 5 11132 0 aVF
+08704_hr.dat 16 1000.0(0)/mV 16 0 210 35283 0 V1
+08704_hr.dat 16 1000.0(0)/mV 16 0 -605 1875 0 V2
+08704_hr.dat 16 1000.0(0)/mV 16 0 -360 20664 0 V3
+08704_hr.dat 16 1000.0(0)/mV 16 0 -255 26244 0 V4
+08704_hr.dat 16 1000.0(0)/mV 16 0 -230 25112 0 V5
+08704_hr.dat 16 1000.0(0)/mV 16 0 60 30065 0 V6

camel_inference/demo/12585_hr.dat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1dbc29ba7d9857a606c23d1a661dee034285f471908206e5428c9b6b7c8ff467
+size 120000

camel_inference/demo/12585_hr.hea

@@ -0,0 +1,13 @@
+12585_hr 12 500 5000
+12585_hr.dat 16 1000.0(0)/mV 16 0 540 5223 0 I
+12585_hr.dat 16 1000.0(0)/mV 16 0 40 11911 0 II
+12585_hr.dat 16 1000.0(0)/mV 16 0 -500 6676 0 III
+12585_hr.dat 16 1000.0(0)/mV 16 0 -290 57071 0 aVR
+12585_hr.dat 16 1000.0(0)/mV 16 0 520 65066 0 aVL
+12585_hr.dat 16 1000.0(0)/mV 16 0 -230 9479 0 aVF
+12585_hr.dat 16 1000.0(0)/mV 16 0 -130 26077 0 V1
+12585_hr.dat 16 1000.0(0)/mV 16 0 -145 27523 0 V2
+12585_hr.dat 16 1000.0(0)/mV 16 0 -155 35476 0 V3
+12585_hr.dat 16 1000.0(0)/mV 16 0 -350 16663 0 V4
+12585_hr.dat 16 1000.0(0)/mV 16 0 -545 60445 0 V5
+12585_hr.dat 16 1000.0(0)/mV 16 0 -105 58137 0 V6

camel_inference/demo/12646_hr.dat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7cabc7399daa14ebf97cc705b9da09aeb3daee12939ca797d145237cdfcf4ed8
+size 120000

camel_inference/demo/12646_hr.hea

@@ -0,0 +1,13 @@
+12646_hr 12 500 5000
+12646_hr.dat 16 1000.0(0)/mV 16 0 -260 38739 0 I
+12646_hr.dat 16 1000.0(0)/mV 16 0 -465 47305 0 II
+12646_hr.dat 16 1000.0(0)/mV 16 0 -205 8550 0 III
+12646_hr.dat 16 1000.0(0)/mV 16 0 363 22259 0 aVR
+12646_hr.dat 16 1000.0(0)/mV 16 0 -27 47829 0 aVL
+12646_hr.dat 16 1000.0(0)/mV 16 0 -335 60490 0 aVF
+12646_hr.dat 16 1000.0(0)/mV 16 0 220 55705 0 V1
+12646_hr.dat 16 1000.0(0)/mV 16 0 120 51889 0 V2
+12646_hr.dat 16 1000.0(0)/mV 16 0 -390 40903 0 V3
+12646_hr.dat 16 1000.0(0)/mV 16 0 -660 39373 0 V4
+12646_hr.dat 16 1000.0(0)/mV 16 0 -770 43235 0 V5
+12646_hr.dat 16 1000.0(0)/mV 16 0 135 65290 0 V6

camel_inference/demo/example_prompt.json

@@ -0,0 +1,8 @@
+[
+  { "type": "text", "text": "You are given two ECG waveforms of a same patient from two different time points."},
+  { "type": "text", "text": "This is the first ECG."},
+  { "type": "ecg", "ecg": "demo/12585_hr" },
+  { "type": "text", "text": "This is the first ECG." },
+  { "type": "ecg", "ecg": "demo/12646_hr" },
+  { "type": "text", "text": "Has non-diagnostic t abnormalities been resolved in the recent tracing compared to the previous one?"}
+]

camel_inference/pyproject.toml

@@ -0,0 +1,36 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools]
+py-modules = ["run_camel"]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools.package-data]
+camel = ["model_registry.yaml"]
+
+[project]
+name = "camel-inference"
+version = "0.1.0"
+description = "Inference-only CLI for CAMEL ECG-language checkpoints"
+readme = "README.md"
+requires-python = ">=3.9"
+authors = [
+  { name = "CAMEL contributors" }
+]
+dependencies = [
+  "numpy",
+  "scipy",
+  "pyyaml",
+  "torch",
+  "transformers",
+  "peft",
+  "accelerate",
+  "sentencepiece",
+  "protobuf"
+]
+
+[project.scripts]
+camel-infer = "run_camel:main"

camel_inference/run_camel.py

@@ -0,0 +1,49 @@
+import argparse
+from camel.camel_model import CAMEL
+
+def main():
+    parser = argparse.ArgumentParser(description="CAMEL")
+    parser.add_argument("--mode", type=str, choices=['forecast', 'base', 'ecgbench'], default='base')
+    parser.add_argument("--device", type=str, default='cuda:0')
+    parser.add_argument("--json", type=str, default=None)
+    parser.add_argument("--text", type=str, default=None)
+    parser.add_argument("--ecgs", type=str, default=None, nargs='+')
+    parser.add_argument("--ecg-configs", type=str, default=None, nargs='+')
+    parser.add_argument("--temperature", type=float, default=0.0)
+    parser.add_argument(
+        "--top-k",
+        dest="top_k",
+        type=int,
+        default=64,
+        help="Top-k sampling cutoff (set <=0 to disable).",
+    )
+    parser.add_argument(
+        "--top-p",
+        dest="top_p",
+        type=float,
+        default=0.95,
+        help="Nucleus sampling cumulative probability cutoff.",
+    )
+    parser.add_argument(
+        "--min-p",
+        dest="min_p",
+        type=float,
+        default=0.0,
+        help="Minimum per-token probability threshold applied after temperature scaling.",
+    )
+    parser.add_argument(
+        "--max-new-tokens",
+        type=int,
+        default=512,
+        help="Maximum number of tokens to generate per sample.",
+    )
+    args = parser.parse_args()
+    
+    model = CAMEL(mode=args.mode, device=args.device)
+    output, prompt = model.run(args)
+    
+    print(f'Prompt: {prompt}')
+    print(f'Prediction: {output}')
+
+if __name__ == "__main__":
+    main()

camel_inference/scripts/download_checkpoints.sh

@@ -0,0 +1,37 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+echo "Installing huggingface_hub if needed..."
+python3 -m pip install -q --user huggingface_hub
+
+echo "Downloading CAMEL checkpoints from Hugging Face..."
+mkdir -p checkpoints
+
+python3 - <<'PY'
+import os, shutil
+from huggingface_hub import hf_hub_download
+
+repo = "CAMEL-ECG/CAMEL"
+files = [
+    "camel_base.pt",
+    "camel_ecginstruct.pt",
+    "camel_forecast.pt"
+]
+
+os.makedirs("checkpoints", exist_ok=True)
+
+for f in files:
+    print(f"Downloading {f}...")
+    src = hf_hub_download(
+        repo_id=repo,
+        filename=f,
+        repo_type="model"
+    )
+    dst = os.path.join("checkpoints", f)
+    shutil.copy2(src, dst)
+    print(f"Saved to {dst}")
+
+print("All checkpoints downloaded.")
+PY
+
+echo "Done."

camel_inference/src/camel/assertions.py

@@ -0,0 +1,509 @@
+"""
+Assertions and summaries for ECG language-model wrappers and their LoRA adapters.
+"""
+from __future__ import annotations
+
+from typing import Any, Dict, Iterable, List, Mapping, Set
+import functools
+import os
+import torch
+
+from camel.ecg_attention_masks import ECGBlockLayout
+
+_ASSERTIONS_ENABLED = os.getenv("ASSERTIONS") == "1"
+
+def _skip_if_assertions_disabled(func):
+    """Decorator that no-ops assertion helpers when ASSERTIONS env var is not set."""
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        if not _ASSERTIONS_ENABLED:
+            return None
+        return func(*args, **kwargs)
+
+    return wrapper
+
+@_skip_if_assertions_disabled
+def assert_tensor_dtype(tensor: torch.Tensor, *, expected: torch.dtype, context: str) -> None:
+    """Verify tensor dtype matches expectation."""
+    if tensor.dtype != expected:
+        raise AssertionError(f"{context}: expected dtype {expected}, got {tensor.dtype}")
+
+@_skip_if_assertions_disabled
+def assert_ecg_blocks_consistent(
+    *,
+    turn_parts: Iterable[Iterable[Dict[str, Any]]],
+    ecg_blocks: Iterable[Dict[str, Any]],
+) -> None:
+    """
+    Validate that structured turn parts contain the expected ECG markers per block.
+    """
+    blocks_list = list(ecg_blocks)
+    expected_counts: Dict[int, Dict[str, int]] = {}
+    per_lead_special_counts: Dict[int, Dict[str, int]] = {}
+    per_lead_secs: Dict[int, Dict[str, Set[int]]] = {}
+    global_special_count: Dict[int, int] = {}
+
+    for idx, blk in enumerate(blocks_list):
+        leads = [str(ld) for ld in blk.get("lead_names", [])]
+        segs = [int(n) for n in blk.get("segments_per_lead", [])]
+        expected_counts[idx] = {ld: int(n) for ld, n in zip(leads, segs)}
+        per_lead_special_counts[idx] = {ld: 0 for ld in leads}
+        per_lead_secs[idx] = {ld: set() for ld in leads}
+        global_special_count[idx] = 0
+
+    for turn in turn_parts:
+        for part in turn:
+            kind = part.get("kind")
+            if kind == "text":
+                continue
+            block_idx = part.get("block_index")
+            if block_idx is None or int(block_idx) not in expected_counts:
+                raise AssertionError("ECG part references unknown block_index.")
+            block_idx = int(block_idx)
+            allowed_leads = set(expected_counts[block_idx].keys())
+            if kind == "special":
+                lead = part.get("lead")
+                if lead:
+                    if lead not in allowed_leads:
+                        raise AssertionError(f"Special token references unknown lead '{lead}'.")
+                    per_lead_special_counts[block_idx][lead] = per_lead_special_counts[block_idx].get(lead, 0) + 1
+                else:
+                    global_special_count[block_idx] = global_special_count.get(block_idx, 0) + 1
+                token_literal = part.get("token")
+                if not isinstance(token_literal, str) or len(token_literal) == 0:
+                    raise AssertionError("Special turn part lacks a string token literal.")
+                continue
+            if kind == "ecg":
+                lead = part.get("lead")
+                if lead not in allowed_leads:
+                    raise AssertionError(f"ECG segment references unknown lead '{lead}'.")
+                sec_val = part.get("sec")
+                expected_total = expected_counts[block_idx].get(lead, 0)
+                if expected_total <= 0:
+                    raise AssertionError(
+                        f"Lead '{lead}' has non-positive declared segments ({expected_total}) but ECG markers are present."
+                    )
+                try:
+                    sec = int(sec_val)
+                except Exception as exc:  # noqa: BLE001
+                    raise AssertionError(f"ECG segment for lead '{lead}' has non-integer sec {sec_val!r}.") from exc
+                if sec < 1 or sec > expected_total:
+                    raise AssertionError(
+                        f"ECG segment for lead '{lead}' has second={sec}, expected within [1,{expected_total}]."
+                    )
+                if sec in per_lead_secs[block_idx][lead]:
+                    raise AssertionError(f"Duplicate ECG segment marker for lead '{lead}' second {sec}.")
+                per_lead_secs[block_idx][lead].add(sec)
+                continue
+            raise AssertionError(f"Unknown turn_parts kind '{kind}'.")
+
+    for block_idx, expected in expected_counts.items():
+        if global_special_count.get(block_idx, 0) != 2:
+            raise AssertionError(
+                f"Expected exactly two global ECG markers for block {block_idx}; "
+                f"found {global_special_count.get(block_idx, 0)}."
+            )
+        for lead, expected_total in expected.items():
+            expected_specials = per_lead_special_counts[block_idx].get(lead, 0)
+            if expected_specials != 2:
+                raise AssertionError(
+                    f"Lead '{lead}' has {expected_specials} special markers; expected start and end (2 total)."
+                )
+            seen_secs = per_lead_secs[block_idx].get(lead, set())
+            if expected_total != len(seen_secs):
+                missing = sorted(set(range(1, expected_total + 1)) - seen_secs)
+                raise AssertionError(
+                    f"Lead '{lead}' missing ECG segment markers for seconds {missing} (expected {expected_total})."
+                )
+
+
+# ---------------- Trainer batch validation helpers -----------------------------------------------
+
+@_skip_if_assertions_disabled
+def assert_prefix_split_complete(*, offset: int, total_prefix_rows: int) -> None:
+    """Validate that prefix splitting consumed all rows."""
+    if offset != total_prefix_rows:
+        raise RuntimeError(
+            f"Prefix split mismatch: consumed {offset} rows but have {total_prefix_rows}"
+        )
+
+
+@_skip_if_assertions_disabled
+def assert_prefix_matches_segments(
+    *,
+    prefix_rows: int,
+    segments_per_lead: Iterable[int],
+    lead_names: Iterable[str],
+    sample_index: int,
+    block_index: int,
+) -> None:
+    """Validate that prefix row count matches sum of segments_per_lead."""
+    total_segments = sum(int(n) for n in segments_per_lead)
+    if prefix_rows != total_segments:
+        raise RuntimeError(
+            f"Sample {sample_index} block {block_index}: Prefix rows ({prefix_rows}) "
+            f"!= sum(segments_per_lead) ({total_segments}). "
+            f"lead_names={list(lead_names)} segments_per_lead={list(segments_per_lead)}"
+        )
+
+@_skip_if_assertions_disabled
+def assert_ecg_part_bounds(
+    *,
+    lead: str,
+    sec: int,
+    lead_to_offset: Mapping[str, int],
+    declared_segments: Mapping[str, int],
+    total_prefix_rows: int,
+    sample_index: int,
+    block_index: int,
+) -> None:
+    """Validate ECG part (lead, sec) falls within expected bounds."""
+    if lead not in declared_segments:
+        raise RuntimeError(f"Unknown lead {lead} in parts for sample {sample_index} block {block_index}")
+
+    nseg = int(declared_segments[lead])
+    if not (1 <= sec <= nseg):
+        raise RuntimeError(
+            f"sec out of range for lead {lead}: got {sec}, expected 1..{nseg}"
+        )
+
+    base = lead_to_offset[lead]
+    start = base
+    end = base + nseg  # exclusive
+    row_idx = start + (sec - 1)
+
+    # Check both global and per-lead bounds
+    if not (0 <= row_idx < total_prefix_rows):
+        raise RuntimeError(
+            f"Bad (lead,sec)=({lead},{sec}) for sample {sample_index} block {block_index}: "
+            f"row_idx {row_idx} not in [0,{total_prefix_rows})"
+        )
+    if not (start <= row_idx < end):
+        raise RuntimeError(
+            f"(lead,sec)=({lead},{sec}) maps outside this lead block "
+            f"[{start},{end}) (row_idx={row_idx}) for sample {sample_index}"
+        )
+
+
+@_skip_if_assertions_disabled
+def assert_layout_specials_complete(
+    *,
+    block_layout: ECGBlockLayout,
+    lead_names: Iterable[str],
+) -> None:
+    """Validate that layout has complete and ordered special token markers.
+
+    For each declared lead:
+      - Both start and end must be present (or both absent)
+      - If present, start < end
+
+    For global markers:
+      - Both start and end must be present (or both absent)
+      - If present, start < end
+    """
+    # Check per-lead specials
+    for ld in lead_names:
+        s = block_layout.lead_start_idx.get(ld)
+        e = block_layout.lead_end_idx.get(ld)
+        if (s is None) != (e is None):
+            raise RuntimeError(f"Lead {ld} missing start/end special (s={s}, e={e})")
+        if s is not None and not (s < e):
+            raise RuntimeError(f"Lead {ld} specials out of order: start={s}, end={e}")
+
+    # Check global specials
+    if (block_layout.global_start_idx is None) != (block_layout.global_end_idx is None):
+        raise RuntimeError("Global start/end special mismatch")
+    if block_layout.global_start_idx is not None and block_layout.global_end_idx is not None:
+        if not (block_layout.global_start_idx < block_layout.global_end_idx):
+            raise RuntimeError(
+                f"Global specials out of order: start={block_layout.global_start_idx} "
+                f"end={block_layout.global_end_idx}"
+            )
+
+# ---------------- Wrapper embedding validations --------------------------------------------------
+
+@_skip_if_assertions_disabled
+def assert_wrapper_embed_length(
+    *,
+    embeddings: torch.Tensor,
+    ids: List[int],
+    context: str,
+) -> None:
+    """Ensure embedding sequence length matches token count exactly.
+
+    This is a critical invariant check that ensures the 1:1 mapping between
+    input token IDs and output embeddings is preserved.
+
+    Args:
+        embeddings: Output embedding tensor (must be at least 1-D)
+        ids: Input token ID list
+        context: Description of where this check is being performed
+
+    Raises:
+        RuntimeError: if ids is not a list, embeddings are not at least 1-D,
+                     or if the embedding length doesn't match the token count
+    """
+    if not isinstance(ids, list):
+        raise RuntimeError(f"{context}: ids must be a python list of ints")
+    if embeddings.dim() < 1:
+        raise RuntimeError(f"{context}: embeddings must be at least 1-D, got shape {tuple(embeddings.shape)}")
+    if embeddings.size(0) != len(ids):
+        raise RuntimeError(f"{context}: embed length {embeddings.size(0)} != token count {len(ids)}")
+
+@_skip_if_assertions_disabled
+def assert_rest_length_nonnegative(*, rest_length: int) -> None:
+    """Validate that rest token length is non-negative.
+
+    This should never happen in correct code, but catching it early helps
+    identify bugs in label construction logic.
+
+    Args:
+        rest_length: Length of ids_rest list
+
+    Raises:
+        ValueError: if rest_length is negative
+    """
+    if rest_length < 0:
+        raise ValueError("ids_rest length is negative (internal error).")
+
+
+# ---------------- Utility assertion helpers ------------------------------------------------------
+
+@_skip_if_assertions_disabled
+def assert_sorted_non_overlapping_spans(spans: List[tuple[int, int]], length: int, ctx: str) -> None:
+    """Validate that spans are sorted, non-overlapping, and within bounds."""
+    prev_end = -1
+    for i, (s, e) in enumerate(spans):
+        if not (0 <= s <= e <= length):
+            raise AssertionError(f"{ctx}: span {i}={(s,e)} out of bounds for length {length}")
+        if s < prev_end:
+            raise AssertionError(f"{ctx}: spans overlap or not sorted at {i-1},{i}: prev_end={prev_end}, curr_start={s}")
+        prev_end = e
+
+
+@_skip_if_assertions_disabled
+def assert_equal_int(a: int, b: int, msg: str) -> None:
+    """Assert two integers are equal with a descriptive message."""
+    if int(a) != int(b):
+        raise AssertionError(f"{msg}: {a} != {b}")
+
+
+@_skip_if_assertions_disabled
+def assert_positive_int(n: int, msg: str) -> None:
+    """Assert an integer is positive (> 0)."""
+    if int(n) <= 0:
+        raise AssertionError(f"{msg}: expected > 0, got {n}")
+
+
+# ---------------- Schema and catalog validations -------------------------------------------------
+
+@_skip_if_assertions_disabled
+def assert_ecg_catalog_valid(catalog: Any, schema: Any) -> None:
+    """Validate ECG special token catalog for uniqueness and mapping consistency.
+
+    Checks:
+      - All tokens are unique
+      - Every canonical lead has entries in lead_to_indices and lead_to_tokens
+      - Token-to-index mappings are consistent across all structures
+      - Global markers (start/end) are present in the catalog
+    """
+    # Uniqueness
+    if len(set(catalog.tokens)) != len(catalog.tokens):
+        raise AssertionError("ECG special tokens contain duplicates")
+
+    # Per-lead mappings
+    for lead in schema.ecg.canonical_leads:
+        if lead not in catalog.lead_to_indices or lead not in catalog.lead_to_tokens:
+            raise AssertionError(f"Missing lead in catalog: {lead}")
+        for kind in ("start", "end"):
+            tok = catalog.lead_to_tokens[lead][kind]
+            idx = catalog.lead_to_indices[lead][kind]
+            if catalog.tokens[idx] != tok:
+                raise AssertionError(f"Catalog mismatch for {lead}:{kind}: tokens[idx] != tok")
+            if catalog.token_to_index.get(tok, None) != idx:
+                raise AssertionError(f"token_to_index mismatch for {lead}:{kind}")
+
+    # Global markers
+    for tok in (schema.ecg.global_start, schema.ecg.global_end):
+        if tok not in catalog.token_to_index:
+            raise AssertionError(f"Global ECG token missing from catalog: {tok}")
+
+
+# ---------------- Conversation and role validations ----------------------------------------------
+
+@_skip_if_assertions_disabled
+def assert_normalized_role_canonical(role: str, schema: Any) -> None:
+    """Ensure normalized role matches one of the canonical prompt roles."""
+    if role not in (schema.prompt.user_role, schema.prompt.model_role):
+        raise AssertionError(f"Normalized role '{role}' did not resolve to a canonical prompt role")
+
+# ---------------- Tokenization and span validations ----------------------------------------------
+
+@_skip_if_assertions_disabled
+def assert_tokenization_cursor_matches(cursor: int, ids_length: int) -> None:
+    """Ensure cursor tracking matches actual text_ids length."""
+    if cursor != ids_length:
+        raise AssertionError(f"cursor ({cursor}) != len(text_ids) ({ids_length})")
+
+
+@_skip_if_assertions_disabled
+def assert_model_spans_valid(model_spans: List[tuple[int, int]], ids_length: int) -> None:
+    """Validate model spans are sorted, non-overlapping, and at least one exists."""
+    assert_sorted_non_overlapping_spans(model_spans, ids_length, ctx="model_spans_in_text")
+    if len(model_spans) == 0:
+        raise AssertionError("No model spans found in text ids")
+
+
+@_skip_if_assertions_disabled
+def assert_eos_appended(ids: List[int], tokenizer: Any, require_eos: bool) -> None:
+    """Validate EOS token was appended if required."""
+    if require_eos and tokenizer.eos_token_id is not None:
+        if not ids or ids[-1] != tokenizer.eos_token_id:
+            raise AssertionError("Required EOS was not appended at the end of text_ids")
+
+
+# ---------------- u0 parts structure validations -------------------------------------------------
+
+@_skip_if_assertions_disabled
+def assert_turn_parts_structure_valid(
+    parts: List[Dict[str, Any]],
+    ecg_blocks: List[Dict[str, Any]],
+    schema: Any,
+    catalog: Any,
+) -> None:
+    """Validate the complete structure of turn parts for all blocks present."""
+    block_indices = sorted({int(p.get("block_index")) for p in parts if p.get("block_index") is not None})
+    for block_idx in block_indices:
+        if block_idx < 0 or block_idx >= len(ecg_blocks):
+            raise AssertionError(f"Unknown block_index {block_idx} in turn parts")
+        blk = ecg_blocks[block_idx]
+        leads_present = [str(ld) for ld in blk.get("lead_names", [])]
+        segments_per_lead = [int(n) for n in blk.get("segments_per_lead", [])]
+
+        special_tokens = [p.get("token") for p in parts if p.get("kind") == "special" and p.get("block_index") == block_idx]
+        if schema.ecg.global_start not in special_tokens:
+            raise AssertionError(f"Missing global_start special in block {block_idx}")
+        if schema.ecg.global_end not in special_tokens:
+            raise AssertionError(f"Missing global_end special in block {block_idx}")
+
+        idx_global_start = next(
+            (i for i, p in enumerate(parts)
+             if p.get("block_index") == block_idx and p.get("kind") == "special"
+             and p.get("token") == schema.ecg.global_start),
+            None
+        )
+        idx_global_end = next(
+            (i for i, p in enumerate(parts)
+             if p.get("block_index") == block_idx and p.get("kind") == "special"
+             and p.get("token") == schema.ecg.global_end),
+            None
+        )
+        if idx_global_start is None or idx_global_end is None or not (idx_global_start < idx_global_end):
+            raise AssertionError(f"Block {block_idx}: missing or misordered global start/end specials")
+
+        for lead, nseg in zip(leads_present, segments_per_lead):
+            assert_positive_int(nseg, f"segments_per_lead for {lead}")
+            idx_start = next(
+                (i for i, p in enumerate(parts)
+                 if p.get("block_index") == block_idx and p.get("kind") == "special" and p.get("lead") == lead
+                 and p.get("token") == catalog.lead_to_tokens[lead]["start"]),
+                None
+            )
+            idx_end = next(
+                (i for i, p in enumerate(parts)
+                 if p.get("block_index") == block_idx and p.get("kind") == "special" and p.get("lead") == lead
+                 and p.get("token") == catalog.lead_to_tokens[lead]["end"]),
+                None
+            )
+            if idx_start is None or idx_end is None or not (idx_start < idx_end):
+                raise AssertionError(f"Lead {lead}: missing or misordered start/end specials in block {block_idx}")
+            secs = [p["sec"] for p in parts[idx_start+1:idx_end] if p.get("kind") == "ecg" and p.get("lead") == lead]
+            if secs != list(range(1, int(nseg) + 1)):
+                raise AssertionError(f"Lead {lead}: ECG seconds sequence invalid: {secs} vs 1..{nseg}")
+            if parts[idx_start]["token_index"] != catalog.lead_to_indices[lead]["start"]:
+                raise AssertionError(f"Lead {lead}: start token_index mismatch")
+            if parts[idx_end]["token_index"] != catalog.lead_to_indices[lead]["end"]:
+                raise AssertionError(f"Lead {lead}: end token_index mismatch")
+
+
+@_skip_if_assertions_disabled
+def assert_turn_content_ends_with_eot(text_block: str, end_of_turn: str) -> None:
+    """Ensure turn content ends with the provided end-of-turn suffix."""
+    if not text_block.endswith(end_of_turn):
+        raise AssertionError("Turn content must end with end_of_turn suffix")
+
+
+# ---------------- Per-sample packing validations -------------------------------------------------
+
+@_skip_if_assertions_disabled
+def assert_leads_canonical_and_ordered(leads_present: List[str], canonical_leads: tuple) -> None:
+    """Validate all leads are canonical (order is explicit and not enforced here)."""
+    lead_list = list(leads_present)
+    if any(ld not in canonical_leads for ld in lead_list):
+        raise AssertionError(f"Non-canonical lead found in leads_present: {lead_list}")
+    if len(set(lead_list)) != len(lead_list):
+        raise AssertionError(f"Duplicate lead detected in leads_present: {lead_list}")
+
+
+@_skip_if_assertions_disabled
+def assert_waveform_shapes_valid(
+    leads_present: List[str],
+    segments_per_lead: List[int],
+    waveform_segments: Dict[str, Any],
+) -> None:
+    """Validate waveform tensor shapes and segment counts.
+
+    Each waveform must be [T, 256] where T matches segments_per_lead.
+    """
+    for ld, nseg in zip(leads_present, segments_per_lead):
+        assert_positive_int(nseg, f"segments_per_lead[{ld}]")
+        wf = waveform_segments[ld]
+        if wf.ndim != 2 or wf.shape[1] != 256:
+            raise AssertionError(f"Waveform for {ld} must be [T,256], got {tuple(wf.shape)}")
+        assert_equal_int(wf.shape[0], nseg, f"Waveform seconds vs segments_per_lead for {ld}")
+
+
+__all__ = [
+    "assertions_active",
+    "capture_adapter_snapshot",
+    "assert_wrapper_adapter_requires_grad",
+    "assert_wrapper_optimizer_coverage",
+    "assert_adapter_gradients",
+    "assert_adapter_updates",
+    "assert_trainable_param_sync",
+    "assert_tensor_dtype",
+    "assert_only_llava_proj_trainable",
+    "summarize_trainables_llava_lora",
+    "assert_language_lora_only",
+    "assert_single_bos_eos",
+    "assert_ecg_layout_valid",
+    "assert_ecg_mask_against_layout",
+    "assert_single_block_mask_matches_reference",
+    "assert_additive_mask_padding",
+    "assert_nonempty_waveform_segments",
+    "assert_prefix_split_complete",
+    "assert_prefix_matches_segments",
+    "assert_ids_are_lists",
+    "assert_embedding_length_matches_tokens",
+    "assert_ecg_part_bounds",
+    "assert_layout_specials_complete",
+    "assert_labels_match_spans",
+    "assert_wrapper_embed_length",
+    "assert_rest_length_nonnegative",
+    "assert_sorted_non_overlapping_spans",
+    "assert_equal_int",
+    "assert_positive_int",
+    "assert_ecg_catalog_valid",
+    "assert_normalized_role_canonical",
+    "assert_rest_blocks_valid",
+    "assert_tokenization_cursor_matches",
+    "assert_model_spans_valid",
+    "assert_turn_parts_consistent",
+    "assert_ecg_blocks_consistent",
+    "assert_eos_appended",
+    "assert_turn_parts_structure_valid",
+    "assert_turn_content_ends_with_eot",
+    "assert_leads_canonical_and_ordered",
+    "assert_waveform_shapes_valid",
+    "assert_collate_item_valid",
+]

camel_inference/src/camel/camel_model.py

@@ -0,0 +1,166 @@
+#!/usr/bin/env python3
+"""Compatibility wrapper around inference.KardiaLM."""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+import json
+import torch
+
+from camel.inference import KardiaLM
+from camel.process_ecg import get_waveform
+
+class CAMEL:
+    def __init__(
+        self,
+        device: torch.device,
+        mode: str,
+        model_config_name: str = 'medgemma-4b-it',
+        conv_ckpt: Optional[str] = None,
+        no_lora: bool = False,
+        mask_strategy: str = 'semantic',
+        **model_args,
+    ) -> None:
+        default_top_k = model_args.pop("default_top_k", 64)
+        default_top_p = float(model_args.pop("default_top_p", 0.95))
+        default_min_p = float(model_args.pop("default_min_p", 0.0))
+
+        # Initialize model
+        if mode == 'base':
+            ckpt = 'checkpoints/camel_base.pt'
+        elif mode == 'ecgbench':
+            ckpt = 'checkpoints/camel_ecginstruct.pt'
+        elif mode == 'forecast':
+            ckpt = 'checkpoints/camel_forecast.pt'
+
+        self.session = KardiaLM(
+            model_registry_path=None,
+            hf_model_id_override=None,
+            model_config_name=model_config_name,
+            adapter_ckpt=ckpt,
+            conv_ckpt=conv_ckpt,
+            no_lora=no_lora,
+            default_max_new_tokens=int(model_args.pop("default_max_new_tokens", 1000)),
+            default_temperature=float(model_args.pop("default_temperature", 1.0)),
+            default_top_k=None if default_top_k is None else int(default_top_k),
+            default_top_p=default_top_p,
+            default_min_p=default_min_p,
+            mask_strategy=mask_strategy,
+            device=device
+        )
+        self.prompt_tokens = self.session.packing_schema.prompt
+        self.device = device
+
+    def run(self, args):
+        if args.json is None and (args.ecgs is None):
+            raise ValueError("Either one of --json or --ecgs should be non-empty.")
+        
+        if args.json is None:
+            text = args.text or ''
+            raw_context = [{'type': 'text', 'text': text}]
+            for ecg in args.ecgs:
+                raw_context.append({'type': 'ecg', 'ecg': ecg})
+        else:
+            try:
+                with open(args.json, "r") as f:
+                    raw_context = json.load(f)
+            except:
+                raise ValueError(f'Failed during reading json: {args.json}')
+
+        content = self._build_content(raw_content=raw_context, ecg_configs=args.ecg_configs)
+        generate_kwargs = dict(
+            content=content,
+            max_new_tokens=args.max_new_tokens,
+            temperature=args.temperature,
+            top_k=args.top_k,
+            top_p=args.top_p,
+            min_p=args.min_p,
+        )
+        generated_text = self.generate(**generate_kwargs)
+        return generated_text
+
+    def generate(
+        self,
+        content: list[dict[str, Any]],
+        max_new_tokens: int = 1000,
+        temperature: float = 1.0,
+        top_k: Optional[int] = 64,
+        top_p: float = 0.95,
+        min_p: float = 0.0,
+    ) -> str:
+        text, prompt_preview = self.session.chat(
+            conversation=content,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            min_p=min_p,
+        )
+        return text, prompt_preview
+
+    def _parse_ecg_config(self, ecg_configs: Optional[list[str]], n_ecgs:int) -> tuple:
+        def _parse_single_config(config):
+            start_ind, end_ind, leads = None, None, None
+            for field in config.split(";"):
+                field = field.strip()
+                if not field:
+                    continue
+                if ":" not in field:
+                    raise ValueError(f"Invalid field: {field}. Expected key:value.")
+
+                key, value = field.split(":", 1)
+                key = key.strip().lower()
+                value = value.strip()
+
+                if key == "start":
+                    start_ind = int(value)
+                elif key == "end":
+                    end_ind = int(value)
+                elif key in ("use_leads", "leads"):
+                    leads = [x.strip() for x in value.split(",") if x.strip()]
+                else:
+                    print(f"Ignoring the unknown key: {key}")
+            return start_ind, end_ind, leads
+        
+        if ecg_configs is None:
+            output = [None] * n_ecgs
+            return output, output, output
+        
+        n_configs = len(ecg_configs)
+        if  n_configs!= 1 and n_configs != n_ecgs:
+            raise ValueError(f'Found {n_configs} ECG configs for {n_ecgs} ECG inputs. The number of config should be 1 or match the number of ECGs.')
+        
+        start_inds, end_inds, leads = [], [], []
+        for config in ecg_configs:
+            start_ind, end_ind, lead = _parse_single_config(config)
+            print(f'ECG Config: {start_ind}, {end_ind}, {lead}')
+            start_inds.append(start_ind)
+            end_inds.append(end_ind)
+            leads.append(lead)
+        
+        if n_configs == 1 and n_ecgs > 1:
+            start_inds = start_inds * n_ecgs
+            end_inds = end_inds * n_ecgs
+            leads = leads * n_ecgs
+
+        return start_inds, end_inds, leads
+
+
+    def _build_content(self, *, raw_content: list[dict[str, str]], ecg_configs: Optional[list[str]]) -> list[dict[str, Any]]:
+        n_ecgs = sum([True for c in raw_content if c['type'] == 'ecg'])
+        starts, ends, leads = self._parse_ecg_config(ecg_configs, n_ecgs)
+        ecg_ind = 0
+
+        content: list[dict[str, Any]] = []
+        for c in raw_content:
+            if c['type'] == 'text':
+                content.append({"type": "text", "text": c['text']})
+            elif c['type'] == 'ecg':
+                waveform = get_waveform(ecg_path=c['ecg'], start_sec=starts[ecg_ind], end_sec=ends[ecg_ind], leads=leads[ecg_ind], device=self.device)
+                ecg_ind += 1
+                content.append({"type": "ecg", "waveform_segments": waveform})
+ 
+        conversation = [{"from": self.prompt_tokens.user_role, "content": content}]
+        return conversation
+
+__all__ = ["CAMEL"]

camel_inference/src/camel/checkpoint_utils.py

@@ -0,0 +1,235 @@
+"""
+Checkpoint- and state-management utilities shared by ECG training scripts.
+These helpers were extracted from train_ecg_text.py to keep the training entrypoint
+focused on orchestration while preserving original behaviour.
+"""
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Tuple
+import torch
+import torch.nn as nn
+from torch.distributed.checkpoint.state_dict import (
+    set_model_state_dict,
+    StateDictOptions,
+)
+from torch.distributed.tensor import DTensor
+from peft import (
+    LoraConfig,
+    TaskType,
+    set_peft_model_state_dict,
+)
+
+from camel.training_setup import is_main_process
+
+def _module_has_dtensor_params(mod: nn.Module) -> bool:
+    """
+    Return True if any parameter tensor underlying the module is a DTensor.
+    Under FSDP2 it is typically parameter.data that carries the DTensor type.
+    """
+    for param in mod.parameters(recurse=True):
+        if isinstance(getattr(param, "data", None), DTensor):
+            return True
+    return False
+
+def _extract_projector_name(payload: Dict[str, Any]) -> Optional[str]:
+    """Return the stored projector name if present."""
+    name = payload.get("projector_name")
+    if isinstance(name, str) and name:
+        return name
+    extra = payload.get("extra")
+    if isinstance(extra, dict):
+        extra_name = extra.get("projector_name")
+        if isinstance(extra_name, str) and extra_name:
+            return extra_name
+    return None
+
+def peek_projector_name(path: str) -> Optional[str]:
+    """Load a checkpoint just far enough to read the projector name metadata."""
+    if path is None:
+        return None
+    payload = torch.load(path, map_location="cpu")
+    if not isinstance(payload, dict):
+        raise RuntimeError(f"Checkpoint {path} must be a dict to inspect projector metadata.")
+    return _extract_projector_name(payload)
+
+def extract_lora_config_from_checkpoints(
+    resume_ckpt_path: Optional[str],
+    load_llava_from: Optional[str],
+) -> Optional[Dict[str, Any]]:
+    """Extract the LoRA configuration embedded in checkpoints (resume path preferred)."""
+    def _load_config(path: Optional[str]) -> Optional[Dict[str, Any]]:
+        if not path:
+            return None
+        try:
+            payload = torch.load(path, map_location="cpu")
+        except Exception as exc:
+            raise RuntimeError(
+                f"Failed to load checkpoint '{path}' while extracting LoRA config"
+            ) from exc
+        lora_payload = payload.get("lora")
+        if isinstance(lora_payload, dict) and isinstance(lora_payload.get("config"), dict):
+            cfg = dict(lora_payload["config"])
+            if "use_dora" in cfg:
+                cfg["use_dora"] = bool(cfg["use_dora"])
+            return cfg
+        return None
+    cfg = _load_config(resume_ckpt_path)
+    if cfg is not None:
+        return cfg
+    return _load_config(load_llava_from)
+
+def load_llava_and_lora(
+    wrapper: nn.Module,
+    model: nn.Module,
+    ckpt_path: str,
+    *,
+    expect_lora: bool,
+    load_lora: bool = True,
+    missing_lora_ok: bool = False,
+) -> Tuple[Dict[str, Any], nn.Module, Optional[LoraConfig]]:
+    """
+    Load llava_proj (mandatory), optional conv encoder weights, ECG special-token
+    embeddings, and LoRA adapters from a checkpoint.
+    """
+    payload = torch.load(ckpt_path, map_location="cpu")
+    if not isinstance(payload, dict):
+        raise RuntimeError(f"Checkpoint {ckpt_path} must be a dict, got {type(payload).__name__}")
+    extra_payload = payload.get("extra") or {}
+    if not isinstance(extra_payload, dict):
+        raise RuntimeError(
+            f"Checkpoint {ckpt_path} has non-dict extra payload of type {type(extra_payload).__name__}"
+        )
+    ckpt_projector = _extract_projector_name(payload)
+    wrapper_projector = getattr(wrapper, "projector_name", None)
+    if ckpt_projector is not None and wrapper_projector is not None and wrapper_projector != ckpt_projector:
+        raise RuntimeError(
+            f"Checkpoint {ckpt_path} projector '{ckpt_projector}' does not match wrapper projector '{wrapper_projector}'."
+        )
+    llava_sd = payload.get("llava_proj")
+    if not isinstance(llava_sd, dict):
+        raise RuntimeError(f"Checkpoint {ckpt_path} missing llava_proj state_dict.")
+    if is_main_process():
+        print(f"[load-llava] Loading llava_proj weights from {ckpt_path}", flush=True)
+    any_llava_dt = _module_has_dtensor_params(wrapper.llava_proj)
+    if any_llava_dt:
+        set_model_state_dict(
+            model=wrapper.llava_proj,
+            model_state_dict=llava_sd,
+            options=StateDictOptions(full_state_dict=True, broadcast_from_rank0=True),
+        )
+    else:
+        wrapper.llava_proj.load_state_dict(llava_sd, strict=True)
+    conv_sd = payload.get("conv")
+    conv_expected = bool(extra_payload.get("conv_trainable"))
+    if conv_sd is None:
+        if conv_expected:
+            raise RuntimeError(
+                f"Checkpoint {ckpt_path} indicates conv_trainable=True but conv weights are missing."
+            )
+    else:
+        if not isinstance(conv_sd, dict):
+            raise RuntimeError(
+                f"Checkpoint {ckpt_path} conv payload must be a state_dict, got {type(conv_sd).__name__}"
+            )
+        if is_main_process():
+            print(f"[load-llava] Loading conv encoder weights from {ckpt_path}", flush=True)
+        any_conv_dt = _module_has_dtensor_params(wrapper.enc)
+        if any_conv_dt:
+            set_model_state_dict(
+                model=wrapper.enc,
+                model_state_dict=conv_sd,
+                options=StateDictOptions(full_state_dict=True, broadcast_from_rank0=True),
+            )
+        else:
+            wrapper.enc.load_state_dict(conv_sd, strict=True)
+    ecg_special_sd = payload.get("ecg_special")
+    if not isinstance(ecg_special_sd, dict):
+        raise RuntimeError(f"Checkpoint {ckpt_path} missing ECG special-token embedding state.")
+    if is_main_process():
+        print(f"[load-llava] Loading ECG special-token embedding from {ckpt_path}", flush=True)
+    any_special_dt = _module_has_dtensor_params(wrapper.ecg_special_embed)
+    if any_special_dt:
+        set_model_state_dict(
+            model=wrapper.ecg_special_embed,
+            model_state_dict=ecg_special_sd,
+            options=StateDictOptions(full_state_dict=True, broadcast_from_rank0=True),
+        )
+    else:
+        wrapper.ecg_special_embed.load_state_dict(ecg_special_sd, strict=True)
+    lora_payload = payload.get("lora")
+    loaded_lora = False
+    created_cfg: Optional[LoraConfig] = None
+    if load_lora and lora_payload is not None:
+        if not isinstance(lora_payload, dict):
+            raise RuntimeError(
+                f"Checkpoint {ckpt_path} has non-dict LoRA payload of type {type(lora_payload).__name__}"
+            )
+        lora_state = lora_payload.get("state_dict")
+        if not isinstance(lora_state, dict):
+            raise RuntimeError(f"Checkpoint {ckpt_path} LoRA payload missing state_dict.")
+        if is_main_process():
+            print(f"[load-llava] Loading LoRA adapters from {ckpt_path}", flush=True)
+        set_peft_model_state_dict(model, lora_state)
+        loaded_lora = True
+        cfg_dict = lora_payload.get("config")
+        if isinstance(cfg_dict, dict):
+            cfg_args = dict(cfg_dict)
+            task_type_raw = cfg_args.get("task_type", TaskType.CAUSAL_LM)
+            if not isinstance(task_type_raw, TaskType):
+                try:
+                    task_type_raw = TaskType(task_type_raw)
+                except Exception:
+                    task_type_raw = TaskType.CAUSAL_LM
+            cfg_args["task_type"] = task_type_raw
+            if "lora_dropout" in cfg_args:
+                try:
+                    cfg_args["lora_dropout"] = float(cfg_args["lora_dropout"])
+                except Exception:
+                    cfg_args["lora_dropout"] = 0.0
+            if "r" in cfg_args:
+                try:
+                    cfg_args["r"] = int(cfg_args["r"])
+                except Exception:
+                    cfg_args["r"] = 0
+            if "lora_alpha" in cfg_args:
+                try:
+                    cfg_args["lora_alpha"] = int(cfg_args["lora_alpha"])
+                except Exception:
+                    cfg_args["lora_alpha"] = 0
+            if "target_modules" in cfg_args and cfg_args["target_modules"] is not None:
+                cfg_args["target_modules"] = list(cfg_args["target_modules"])
+            cfg_args.setdefault("bias", "none")
+            cfg_args.setdefault("inference_mode", False)
+            cfg_args["use_dora"] = bool(cfg_args.get("use_dora", False))
+            try:
+                created_cfg = LoraConfig(**cfg_args)
+            except Exception:
+                created_cfg = None
+    if expect_lora and load_lora and not loaded_lora:
+        if missing_lora_ok:
+            if is_main_process():
+                print(
+                    f"[load-llava] Warning: checkpoint {ckpt_path} contains no LoRA adapters; "
+                    "continuing with the currently-initialized adapters.",
+                    flush=True,
+                )
+        else:
+            raise RuntimeError(
+                f"[load-llava] Warning: expected LoRA adapters in {ckpt_path} but none were loaded."
+            )
+    return extra_payload, model, created_cfg
+
+def update_wrapper_language_model(wrapper: nn.Module, model: nn.Module) -> None:
+    """Ensure the wrapper references the latest language-model instance."""
+    wrapper.language_model = model
+
+__all__ = [
+    "find_latest_step_checkpoint",
+    "extract_lora_config_from_checkpoints",
+    "dump_lora_state_fsdp_safe",
+    "prepare_optimizer_state_payload",
+    "load_llava_and_lora",
+    "update_wrapper_language_model",
+    "ensure_no_dtensor",
+    "peek_projector_name",
+]

camel_inference/src/camel/ecg_attention_masks.py

@@ -0,0 +1,343 @@
+"""ECG-aware attention mask utilities with pluggable strategy support."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Protocol
+import torch
+
+@dataclass
+class ECGBlockLayout:
+    """Layout describing a single ECG block inside the assembled sequence."""
+
+    start_idx: Optional[int]
+    end_idx_exclusive: Optional[int]
+    global_start_idx: Optional[int] = None
+    global_end_idx: Optional[int] = None
+    lead_start_idx: Dict[str, int] = field(default_factory=dict)
+    lead_end_idx: Dict[str, int] = field(default_factory=dict)
+    signal_pos_by_lead: Dict[str, List[int]] = field(default_factory=dict)
+    time_to_signal_idxs: Dict[int, List[int]] = field(default_factory=dict)
+    special_idxs_sorted: List[int] = field(default_factory=list)
+    signal_pos_list: List[int] = field(default_factory=list)
+    declared_segments_per_lead: Dict[str, int] = field(default_factory=dict)
+    conv_idxs: List[int] = field(default_factory=list)
+
+
+@dataclass
+class ECGSequenceLayout:
+    """Compact description of the assembled token layout for one training sample."""
+
+    seq_len: int
+    text_idxs: List[int] = field(default_factory=list)
+    blocks: List[ECGBlockLayout] = field(default_factory=list)
+
+
+def _as_tensor(indices: List[int], device: torch.device) -> torch.Tensor:
+    if not indices:
+        return torch.empty(0, dtype=torch.long, device=device)
+    return torch.tensor(sorted(indices), dtype=torch.long, device=device)
+
+
+def apply_single_block_semantic_mask_(
+    allowed: torch.Tensor,
+    block_layout: ECGBlockLayout,
+    *,
+    visible_prefix_len: int,
+    key_limit_exclusive: int,
+    apply_header_causal: bool = True,
+) -> None:
+    """
+    In-place semantic mask update for a single ECG block. Mirrors the historical
+    single-block logic, but operates on a provided boolean mask.
+    """
+    L = int(allowed.size(0))
+    device = allowed.device
+    header = (
+        torch.arange(int(visible_prefix_len), dtype=torch.long, device=device)
+        if int(visible_prefix_len) > 0
+        else torch.empty(0, dtype=torch.long, device=device)
+    )
+
+    specials_list = block_layout.special_idxs_sorted or []
+    if not specials_list:
+        specials_list = sorted(
+            ([block_layout.global_start_idx] if block_layout.global_start_idx is not None else [])
+            + list(block_layout.lead_start_idx.values())
+            + list(block_layout.lead_end_idx.values())
+            + ([block_layout.global_end_idx] if block_layout.global_end_idx is not None else [])
+        )
+        block_layout.special_idxs_sorted = specials_list
+    if key_limit_exclusive is not None:
+        specials_list = [i for i in specials_list if int(i) < int(key_limit_exclusive)]
+    specials = _as_tensor([int(i) for i in specials_list], device)
+
+    signals_list = block_layout.signal_pos_list or []
+    if not signals_list:
+        signal_all: List[int] = []
+        for lst in block_layout.signal_pos_by_lead.values():
+            signal_all.extend(lst)
+        signals_list = sorted(signal_all)
+        block_layout.signal_pos_list = signals_list
+    if key_limit_exclusive is not None:
+        signals_list = [i for i in signals_list if int(i) < int(key_limit_exclusive)]
+    signals = _as_tensor([int(i) for i in signals_list], device)
+
+    lead_starts = _as_tensor(list(block_layout.lead_start_idx.values()), device)
+    lead_ends = _as_tensor(list(block_layout.lead_end_idx.values()), device)
+
+    if apply_header_causal and header.numel():
+        allowed[header[:, None], header[None, :]] = header[:, None] >= header[None, :]
+
+    gs = block_layout.global_start_idx
+    if gs is not None and header.numel():
+        allowed[int(gs), header] = True
+
+    rows_before = []
+    if lead_starts.numel():
+        rows_before.append(lead_starts)
+    if signals.numel():
+        rows_before.append(signals)
+    if lead_ends.numel():
+        rows_before.append(lead_ends)
+    rows_before_t = (
+        torch.cat(rows_before, dim=0) if rows_before else torch.empty(0, dtype=torch.long, device=device)
+    )
+
+    if rows_before_t.numel():
+        if header.numel():
+            allowed[rows_before_t[:, None], header[None, :]] = True
+        if specials.numel():
+            allowed[rows_before_t[:, None], specials[None, :]] = specials[None, :] < rows_before_t[:, None]
+
+    ttsi: Dict[int, List[int]] = block_layout.time_to_signal_idxs
+    if signals.numel() and ttsi:
+        pos_min_time: Dict[int, int] = {}
+        pos_to_time: Dict[int, int] = {}
+        for t, idxs in ttsi.items():
+            for p in idxs:
+                pos_to_time[p] = t
+                prev = pos_min_time.get(p)
+                if prev is None or t < prev:
+                    pos_min_time[p] = t
+
+        u_pos_list = sorted(pos_min_time.keys())
+        if u_pos_list:
+            u_pos = torch.tensor(u_pos_list, dtype=torch.long, device=device)
+            u_time = torch.tensor([pos_min_time[p] for p in u_pos_list], dtype=torch.long, device=device)
+            q_time = torch.tensor([pos_to_time.get(p, 0) for p in signals_list], dtype=torch.long, device=device)
+            allowed[signals[:, None], u_pos[None, :]] = (u_time[None, :] <= q_time[:, None])
+
+    for lead, eidx in block_layout.lead_end_idx.items():
+        lead_sigs = block_layout.signal_pos_by_lead.get(lead, [])
+        if lead_sigs:
+            allowed[int(eidx), torch.tensor(lead_sigs, dtype=torch.long, device=device)] = True
+
+    ge = block_layout.global_end_idx
+    if ge is not None:
+        gei = int(ge)
+        if header.numel():
+            allowed[gei, header] = True
+        if specials.numel():
+            allowed[gei, specials] = True
+        if signals.numel():
+            allowed[gei, signals] = True
+
+    conv = _as_tensor(block_layout.conv_idxs, device)
+    if conv.numel():
+        allowed[conv[:, None], conv[None, :]] = conv[:, None] >= conv[None, :]
+        if header.numel():
+            allowed[conv[:, None], header[None, :]] = True
+        if specials.numel():
+            allowed[conv[:, None], specials[None, :]] = True
+        if signals.numel():
+            allowed[conv[:, None], signals[None, :]] = True
+        cols = torch.arange(L, device=device)
+        conv_rows = allowed[conv, :]
+        conv_rows &= (cols.unsqueeze(0) <= conv.unsqueeze(1))
+        allowed[conv, :] = conv_rows
+
+    if specials.numel():
+        allowed[specials, specials] = True
+
+    if key_limit_exclusive is not None and int(key_limit_exclusive) < L:
+        block_rows_list = list(specials_list) + list(signals_list)
+        if block_rows_list:
+            block_rows = _as_tensor(block_rows_list, device)
+            allowed[block_rows, int(key_limit_exclusive):] = False
+
+@dataclass
+class MaskBuildResult:
+    """Container for per-sample mask artifacts produced by a strategy."""
+
+    additive: torch.Tensor
+    boolean: Optional[torch.Tensor] = None
+
+
+class ECGMaskStrategy(Protocol):
+    """Protocol for strategies that build and update per-sample attention masks."""
+
+    name: str
+
+    def build(
+        self,
+        layout: ECGSequenceLayout,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> MaskBuildResult:
+        ...
+
+    def update_for_generated_token(
+        self,
+        layout: ECGSequenceLayout,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        previous: MaskBuildResult,
+    ) -> MaskBuildResult:
+        ...
+
+
+class SemanticMaskStrategy:
+    """Default strategy reproducing the historical ECG-aware attention mask."""
+
+    name = "semantic"
+
+    def build(
+        self,
+        layout: ECGSequenceLayout,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> MaskBuildResult:
+        L = int(layout.seq_len)
+        if L <= 0:
+            raise ValueError("Semantic mask requires a positive sequence length")
+        allowed = torch.tril(torch.ones((L, L), dtype=torch.bool, device=device))
+        multi_block = len(layout.blocks) > 1
+        for block in layout.blocks:
+            if block.start_idx is None or block.end_idx_exclusive is None:
+                continue
+            block_rows: List[int] = []
+            specials_list = block.special_idxs_sorted or []
+            if not specials_list:
+                specials_list = sorted(
+                    ([block.global_start_idx] if block.global_start_idx is not None else [])
+                    + list(block.lead_start_idx.values())
+                    + list(block.lead_end_idx.values())
+                    + ([block.global_end_idx] if block.global_end_idx is not None else [])
+                )
+                block.special_idxs_sorted = specials_list
+            signals_list = block.signal_pos_list or []
+            if not signals_list:
+                signal_all: List[int] = []
+                for lst in block.signal_pos_by_lead.values():
+                    signal_all.extend(lst)
+                signals_list = sorted(signal_all)
+                block.signal_pos_list = signals_list
+            if specials_list:
+                block_rows.extend(int(i) for i in specials_list)
+            if signals_list:
+                block_rows.extend(int(i) for i in signals_list)
+            if block_rows:
+                rows = torch.tensor(sorted(set(block_rows)), dtype=torch.long, device=device)
+                allowed[rows, :] = False
+            apply_single_block_semantic_mask_(
+                allowed,
+                block,
+                visible_prefix_len=int(block.start_idx),
+                key_limit_exclusive=int(block.end_idx_exclusive),
+                apply_header_causal=not multi_block,
+            )
+        additive = self._boolean_to_additive(allowed, device=device, dtype=dtype)
+        return MaskBuildResult(additive=additive, boolean=allowed)
+
+    def update_for_generated_token(
+        self,
+        layout: ECGSequenceLayout,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        previous: MaskBuildResult,
+    ) -> MaskBuildResult:
+        allowed = previous.boolean
+        if allowed is None:
+            return self.build(layout, device=device, dtype=dtype)
+        prev_len = int(allowed.size(0))
+        new_allowed = torch.zeros((prev_len + 1, prev_len + 1), dtype=torch.bool, device=device)
+        new_allowed[:prev_len, :prev_len] = allowed
+        new_allowed[prev_len, : prev_len + 1] = True
+        additive = self._boolean_to_additive(new_allowed, device=device, dtype=dtype)
+        return MaskBuildResult(additive=additive, boolean=new_allowed)
+
+
+    @classmethod
+    def _build_boolean_mask(cls, layout: "ECGSequenceLayout", device: torch.device) -> torch.Tensor:
+        if len(layout.blocks) != 1:
+            raise ValueError("Single-block mask builder requires exactly one ECG block")
+        L = int(layout.seq_len)
+        allowed = torch.zeros((L, L), dtype=torch.bool, device=device)
+        block = layout.blocks[0]
+        prefix_len = int(block.start_idx or 0)
+        end_idx = int(block.end_idx_exclusive or L)
+        conv_idxs = [idx for idx in layout.text_idxs if int(idx) >= end_idx]
+        block_ref = ECGBlockLayout(
+            start_idx=block.start_idx,
+            end_idx_exclusive=block.end_idx_exclusive,
+            global_start_idx=block.global_start_idx,
+            global_end_idx=block.global_end_idx,
+            lead_start_idx=dict(block.lead_start_idx),
+            lead_end_idx=dict(block.lead_end_idx),
+            signal_pos_by_lead=dict(block.signal_pos_by_lead),
+            time_to_signal_idxs=dict(block.time_to_signal_idxs),
+            special_idxs_sorted=list(block.special_idxs_sorted),
+            signal_pos_list=list(block.signal_pos_list),
+            declared_segments_per_lead=dict(block.declared_segments_per_lead),
+            conv_idxs=conv_idxs,
+        )
+        apply_single_block_semantic_mask_(
+            allowed,
+            block_ref,
+            visible_prefix_len=prefix_len,
+            key_limit_exclusive=end_idx,
+        )
+        return allowed
+
+    @staticmethod
+    def _boolean_to_additive(allowed: torch.Tensor, *, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+        additive = torch.zeros(allowed.shape, dtype=dtype, device=device)
+        additive.masked_fill_(~allowed, float("-inf"))
+        return additive
+
+
+DEFAULT_MASK_STRATEGY = SemanticMaskStrategy()
+
+
+MASK_STRATEGY_REGISTRY: Dict[str, ECGMaskStrategy] = {
+    DEFAULT_MASK_STRATEGY.name: DEFAULT_MASK_STRATEGY,
+}
+
+
+def get_mask_strategy(name: Optional[str]) -> ECGMaskStrategy:
+    """Resolve a registered mask strategy by name (case-insensitive)."""
+    if name is None:
+        return DEFAULT_MASK_STRATEGY
+    key = str(name).lower()
+    try:
+        return MASK_STRATEGY_REGISTRY[key]
+    except KeyError as exc:
+        known = ", ".join(sorted(MASK_STRATEGY_REGISTRY))
+        raise ValueError(f"Unknown ECG mask strategy '{name}'. Known strategies: {known}") from exc
+
+
+__all__ = [
+    "ECGBlockLayout",
+    "ECGSequenceLayout",
+    "MaskBuildResult",
+    "ECGMaskStrategy",
+    "SemanticMaskStrategy",
+    "DEFAULT_MASK_STRATEGY",
+    "MASK_STRATEGY_REGISTRY",
+    "get_mask_strategy",
+]

camel_inference/src/camel/ecg_gemma_model.py

@@ -0,0 +1,91 @@
+# ecg_gemma_model.py
+from __future__ import annotations
+
+from typing import Optional
+import torch
+import torch.nn as nn
+from torch import Tensor
+from transformers import AutoModelForCausalLM
+
+from camel.ecg_model_wrapper import ECGLanguageModelWrapper
+
+class ECGGemmaPrefix(ECGLanguageModelWrapper):
+    """
+    Frozen: Gemma-IT (language model), 1D conv signal encoder (loaded from disk)
+    Trainable: llava_proj (Linear 64 -> Gemma hidden size)
+    Optionally trainable: conv encoder when explicitly requested
+
+    This wrapper turns per-second ECG windows into single "pseudo-token" rows
+    that are interleaved into user turns at embedding time.
+    """
+
+    def __init__(
+        self,
+        gemma: AutoModelForCausalLM,
+        enc: nn.Module,
+        hidden_size: int,
+        num_ecg_special_tokens: int,
+        dtype: Optional[torch.dtype] = torch.bfloat16,
+        enc_out_dim: int = 64,  # from the specified conv stack: 4 channels * 16 length = 64
+        freeze_encoder: bool = True,
+        inference: bool = False,
+        projector_name: str = "linear",
+    ):
+        super().__init__(
+            language_model=gemma,
+            conv_encoder=enc,
+            hidden_size=hidden_size,
+            num_ecg_special_tokens=num_ecg_special_tokens,
+            dtype=dtype,
+            enc_out_dim=enc_out_dim,
+            freeze_encoder=freeze_encoder,
+            inference=inference,
+            projector_name=projector_name,
+        )
+        # Use language_model consistently; no Gemma alias is set.
+
+    def forward_language_model(
+        self,
+        inputs_embeds: Tensor,
+        attention_mask: Tensor,
+        labels: Optional[Tensor],
+        output_hidden_states = False,
+    ):
+        # Ensure inputs live on the device of the LM input embeddings when sharded
+        embedder_fn = getattr(self.language_model, "get_input_embeddings", None)
+        if callable(embedder_fn):
+            try:
+                embed_module = embedder_fn()
+            except Exception as exc:
+                raise RuntimeError("Failed to obtain language-model input embeddings for device inference") from exc
+            if not hasattr(embed_module, "weight"):
+                raise RuntimeError("Input embedding module lacks a weight parameter; cannot infer device")
+            dev0 = embed_module.weight.device
+        else:
+            params_iter = self.language_model.parameters()
+            try:
+                first_param = next(params_iter)
+            except StopIteration as exc:
+                raise RuntimeError("Language model exposes no parameters to infer device placement") from exc
+            dev0 = first_param.device
+
+        if inputs_embeds.device != dev0:
+            inputs_embeds = inputs_embeds.to(dev0)
+        if attention_mask.device != dev0:
+            attention_mask = attention_mask.to(dev0)
+        if labels is not None and labels.device != dev0:
+            labels = labels.to(dev0)
+
+        return self.language_model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            labels=labels,
+            use_cache=False,
+            output_hidden_states=output_hidden_states,
+        )
+
+    
+    # FSDP2 helpers removed (not used)
+
+
+__all__ = ["ECGGemmaPrefix"]

camel_inference/src/camel/ecg_model_wrapper.py

@@ -0,0 +1,394 @@
+"""
+Abstract ECG-language model adapter with shared conv→projection logic.
+
+Subclasses can override only the pieces that differ (e.g., prompt format or
+language-model forwarding) while inheriting the common ECG prefix handling.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, Iterable, List, Mapping, Optional, Tuple
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from camel.assertions import assert_wrapper_embed_length, assert_rest_length_nonnegative, assert_tensor_dtype
+from camel.projectors import build_projector
+
+class ECGNonFiniteInputError(RuntimeError):
+    """Raised when the conv encoder input contains NaN or Inf values."""
+
+    def __init__(self, sample_idx: int, lead: Optional[str] = None) -> None:
+        self.sample_idx = int(sample_idx)
+        self.lead = lead
+        lead_part = f", lead={lead}" if lead is not None else ""
+        super().__init__(f"Non-finite waveform detected (sample_idx={self.sample_idx}{lead_part})")
+
+
+class ECGLanguageModelWrapper(nn.Module):
+    """
+    Turns per-lead ECG waveforms into prefix embeddings consumable by a language model.
+    Stores the ECG encoder, trainable adapter, and the target language model.
+    """
+
+    def __init__(
+        self,
+        *,
+        language_model: nn.Module,
+        conv_encoder: nn.Module,
+        hidden_size: int,
+        num_ecg_special_tokens: int,
+        dtype: Optional[torch.dtype] = torch.bfloat16,
+        enc_out_dim: int = 64,
+        freeze_encoder: bool = True,
+        inference: bool = False,
+        projector_name: str = "linear",
+    ) -> None:
+        super().__init__()
+
+        if int(num_ecg_special_tokens) <= 0:
+            raise ValueError("num_ecg_special_tokens must be positive")
+
+        # Keep LM in train mode (actual freezing handled by caller).
+        self.language_model = language_model.train()
+
+        # Conv encoder may be frozen or trainable depending on configuration.
+        self.enc = conv_encoder
+        if freeze_encoder:
+            self.enc = self.enc.eval()
+            for p in self.enc.parameters():
+                p.requires_grad = False
+        else:
+            self.enc = self.enc.train()
+            for p in self.enc.parameters():
+                p.requires_grad = True
+
+        self.hidden_size = int(hidden_size)
+        self.dtype = dtype or torch.bfloat16
+        self.num_ecg_special_tokens = int(num_ecg_special_tokens)
+        self.inference = bool(inference)
+        self.projector_name = str(projector_name or "linear")
+
+        # Trainable adapter: conv output → LM hidden size (kept in fp32 for stability).
+        self.llava_proj = build_projector(self.projector_name, int(enc_out_dim), self.hidden_size)
+        if self.inference:
+            self.llava_proj.to(dtype=torch.float32)
+
+        self.ecg_special_embed = nn.Embedding(self.num_ecg_special_tokens, self.hidden_size)
+        self.ecg_special_embed.to(dtype=self.dtype)
+        if self.inference:
+            self.enc.to(dtype=torch.float32)
+        self._grad_ckpt_enabled = self._detect_grad_ckpt_state()
+
+    def _projector_param_dtype(self) -> torch.dtype:
+        """Return dtype of the first projector parameter (defaults to fp32)."""
+        first_param = next(self.llava_proj.parameters(), None)
+        if first_param is None:
+            return torch.float32
+        return first_param.dtype
+
+    # ---- Gradient checkpointing --------------------------------------------------------------
+
+    def _detect_grad_ckpt_state(self) -> bool:
+        if hasattr(self.language_model, "is_gradient_checkpointing"):
+            try:
+                return bool(self.language_model.is_gradient_checkpointing)
+            except Exception:
+                return False
+        if hasattr(self.language_model, "gradient_checkpointing"):
+            try:
+                return bool(self.language_model.gradient_checkpointing)
+            except Exception:
+                return False
+        return False
+
+    def set_gradient_checkpointing(self, enabled: bool) -> bool:
+        enabled = bool(enabled)
+        current = self._detect_grad_ckpt_state()
+        if current == enabled:
+            self._grad_ckpt_enabled = enabled
+            return False
+        if enabled:
+            if hasattr(self.language_model, "gradient_checkpointing_enable"):
+                self.language_model.gradient_checkpointing_enable()
+            if hasattr(getattr(self.language_model, "config", None), "use_cache"):
+                self.language_model.config.use_cache = False
+        else:
+            if hasattr(self.language_model, "gradient_checkpointing_disable"):
+                self.language_model.gradient_checkpointing_disable()
+            elif hasattr(self.language_model, "gradient_checkpointing"):
+                try:
+                    self.language_model.gradient_checkpointing = False
+                except Exception:
+                    pass
+            if hasattr(getattr(self.language_model, "config", None), "use_cache"):
+                self.language_model.config.use_cache = True
+        self._grad_ckpt_enabled = enabled
+        return True
+
+    def enable_gradient_checkpointing(self) -> None:
+        self.set_gradient_checkpointing(True)
+
+    def disable_gradient_checkpointing(self) -> None:
+        self.set_gradient_checkpointing(False)
+
+    def is_gradient_checkpointing_enabled(self) -> bool:
+        return bool(self._grad_ckpt_enabled)
+
+    # ---- Token helpers -----------------------------------------------------------------------
+
+    def tokens_to_embeds(self, input_embedder: nn.Embedding, ids: List[int], device: torch.device) -> Tensor:
+        ids_t = torch.tensor(ids, dtype=torch.long, device=device)
+        embeddings = input_embedder(ids_t)
+        embeddings = embeddings.to(dtype=self.dtype)
+        # Defensive: enforce 1:1 mapping between ids and embeddings
+        assert_wrapper_embed_length(embeddings=embeddings, ids=ids, context="tokens_to_embeds")
+        return embeddings
+
+    def ecg_special_tokens_to_embeds(self, indices: torch.Tensor | List[int], device: torch.device) -> Tensor:
+        if torch.is_tensor(indices):
+            idx = indices.to(device=device, dtype=torch.long)
+        else:
+            idx = torch.tensor(indices, dtype=torch.long, device=device)
+        embeds = self.ecg_special_embed(idx)
+        return embeds.to(dtype=self.dtype)
+
+    # ---- ECG prefix encoding -----------------------------------------------------------------
+
+    def ecg_prefix(
+        self,
+        waveform_segments: Dict[str, Tensor],
+        device: torch.device,
+        lead_order: Optional[List[str]] = None,
+    ) -> Tensor:
+        """Encode a single sample's ECG prefix in a deterministic lead order.
+
+        Args:
+            waveform_segments: Mapping lead -> [T,256] windows.
+            device: Target device for encoder input.
+            lead_order: Optional explicit order of leads to iterate. If provided,
+                segments are concatenated in this exact order; otherwise relies on
+                insertion order of the mapping.
+        """
+        seqs: List[Tensor] = []
+        if lead_order is None:
+            items = list(waveform_segments.items())
+        else:
+            items = [(ld, waveform_segments[ld]) for ld in lead_order if ld in waveform_segments]
+            # Validate presence when an explicit order is provided
+            missing = [ld for ld in lead_order if ld not in waveform_segments]
+            if missing:
+                raise ValueError(f"Missing leads in waveform_segments for requested order: {missing}")
+        for lead, seg in items:
+            seg = torch.as_tensor(seg)
+            if seg.ndim != 2 or seg.shape[1] != 256:
+                raise ValueError(f"Waveform for lead {lead} must be [T,256], got {seg.shape}")
+            seqs.append(seg)
+        x = torch.cat(seqs, dim=0)
+        x = x.to(device=device, dtype=torch.float32)
+        x = x.unsqueeze(1)  # [P, 1, 256]
+
+        enc_trainable = any(p.requires_grad for p in self.enc.parameters())
+        ctx = torch.enable_grad() if enc_trainable else torch.no_grad()
+        with ctx:
+            z = self.enc(x)  # [P, C, L]
+        if self.inference:
+            z = z.to(dtype=torch.float32)
+        assert_tensor_dtype(z, expected=torch.float32, context="conv encoder output (single)")
+        self.ensure_finite(z, "conv encoder output")
+        z = z.flatten(1)  # [P, 64] for conv stack
+
+        proj_dtype = self._projector_param_dtype()
+        if z.dtype != proj_dtype:
+            z = z.to(dtype=proj_dtype)
+        y = self.llava_proj(z)
+        if self.inference:
+            y = y.to(dtype=torch.float32)
+        assert_tensor_dtype(y, expected=torch.float32, context="llava_proj output (single)")
+        return y.to(dtype=self.dtype)
+
+    def ecg_prefix_batch(
+        self,
+        waveform_segments_batch: List[Dict[str, Tensor]],
+        device: torch.device,
+        lead_orders: Optional[List[List[str]]] = None,
+    ) -> Tuple[Tensor, List[int]]:
+        """Encode a batch of ECG prefixes with explicit per-sample lead orders.
+
+        Args:
+            waveform_segments_batch: List of dicts; each maps lead -> [T,256].
+            device: Target device for encoder input.
+            lead_orders: Optional list of lead-order lists, one per sample. If
+                provided, each sample's segments are concatenated in that exact
+                order; otherwise relies on insertion order of each mapping.
+        """
+        all_seqs: List[Tensor] = []
+        prefix_lengths: List[int] = []
+
+        for i, wv_dict in enumerate(waveform_segments_batch):
+            seqs: List[Tensor] = []
+            leads_for_sample: List[str] = []
+            if lead_orders is not None and i < len(lead_orders) and lead_orders[i] is not None:
+                order = lead_orders[i]
+                missing = [ld for ld in order if ld not in wv_dict]
+                if missing:
+                    raise ValueError(f"Missing leads for sample {i}: {missing}")
+                items = [(ld, wv_dict[ld]) for ld in order]
+            else:
+                items = list(wv_dict.items())
+            for lead, seg in items:
+                seg = torch.as_tensor(seg)
+                if seg.ndim != 2 or seg.shape[1] != 256:
+                    raise ValueError(f"Waveform for lead {lead} must be [T,256], got {seg.shape}")
+                seqs.append(seg)
+                leads_for_sample.append(str(lead))
+            sample_segments = torch.cat(seqs, dim=0)
+            if not torch.isfinite(sample_segments).all().item():
+                bad_lead = None
+                for lead_name, seg_tensor in zip(leads_for_sample, seqs):
+                    if not torch.isfinite(seg_tensor).all().item():
+                        bad_lead = lead_name
+                        break
+                raise ECGNonFiniteInputError(sample_idx=i, lead=bad_lead)
+            all_seqs.append(sample_segments)
+            prefix_lengths.append(sample_segments.size(0))
+
+        x = torch.cat(all_seqs, dim=0)
+        x = x.to(device=device, dtype=torch.float32).unsqueeze(1)
+
+        enc_trainable = any(p.requires_grad for p in self.enc.parameters())
+        ctx = torch.enable_grad() if enc_trainable else torch.no_grad()
+        with ctx:
+            z = self.enc(x)
+        if self.inference:
+            z = z.to(dtype=torch.float32)
+        assert_tensor_dtype(z, expected=torch.float32, context="conv encoder output (batch)")
+        self.ensure_finite(z, "conv encoder output (batch)")
+        z = z.flatten(1)
+
+        proj_dtype = self._projector_param_dtype()
+        if z.dtype != proj_dtype:
+            z = z.to(dtype=proj_dtype)
+        y = self.llava_proj(z)
+        if self.inference:
+            y = y.to(dtype=torch.float32)
+        assert_tensor_dtype(y, expected=torch.float32, context="llava_proj output (batch)")
+        return y.to(dtype=self.dtype), prefix_lengths
+
+    # ---- Language-model forward --------------------------------------------------------------
+
+    def forward_language_model(
+        self,
+        inputs_embeds: Tensor,
+        attention_mask: Tensor,
+        labels: Optional[Tensor],
+    ):
+        """
+        Default HF-style forward call; subclasses can override for custom behavior.
+        """
+        embedder_fn = getattr(self.language_model, "get_input_embeddings", None)
+        if callable(embedder_fn):
+            embed_module = embedder_fn()
+            if not hasattr(embed_module, "weight"):
+                raise RuntimeError("Language model embeddings missing weight parameter; cannot infer device.")
+            target_device = embed_module.weight.device
+        else:
+            try:
+                first_param = next(self.language_model.parameters())
+            except StopIteration as exc:
+                raise RuntimeError("Language model exposes no parameters to infer device placement.") from exc
+            target_device = first_param.device
+
+        if inputs_embeds.device != target_device:
+            inputs_embeds = inputs_embeds.to(target_device)
+        if attention_mask.device != target_device:
+            attention_mask = attention_mask.to(target_device)
+        if labels is not None and labels.device != target_device:
+            labels = labels.to(target_device)
+
+        return self.language_model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            labels=labels,
+            use_cache=False,
+        )
+
+    # ---- Label helpers ----------------------------------------------------------------------
+
+    @staticmethod
+    def build_labels_from_lengths(
+        *,
+        ids_rest: List[int],
+        model_spans_in_rest: List[Tuple[int, int]],
+        total_len: int,
+        offset_rest: int,
+    ) -> Tensor:
+        """Build label tensor using explicit offsets to satisfy any packing schema.
+
+        Args:
+            ids_rest: Token ids for the supervised span of the prompt.
+            model_spans_in_rest: List of (start, end) spans (relative to ids_rest) to supervise.
+            total_len: Total sequence length of the assembled prompt.
+            offset_rest: Absolute position in the sequence where `ids_rest` begins.
+
+        Returns:
+            Tensor of shape (total_len,) with supervised ids placed according to spans; all other
+            positions are filled with -100.
+
+        Raises:
+            ValueError: if `offset_rest` is invalid or spans fall outside the provided bounds.
+        """
+        labels = torch.full((total_len,), fill_value=-100, dtype=torch.long)
+
+        offset_rest = int(offset_rest)
+        if offset_rest < 0 or offset_rest > total_len:
+            raise ValueError(f"Invalid rest offset {offset_rest} for sequence length {total_len}.")
+        rest_len = len(ids_rest)
+        if offset_rest + rest_len > total_len:
+            raise ValueError(
+                f"Rest tokens (len={rest_len}) exceed total_len {total_len} with offset {offset_rest}."
+            )
+        # Defensive: no negative rest length (should not happen, but keep invariant tight)
+        assert_rest_length_nonnegative(rest_length=rest_len)
+
+        for (s, e) in model_spans_in_rest:
+            if not (0 <= s <= e <= rest_len):
+                raise ValueError(
+                    f"Model span {(s, e)} is out of bounds for ids_rest length {rest_len}."
+                )
+            s_abs = offset_rest + s
+            e_abs = offset_rest + e
+            if not (0 <= s_abs <= e_abs <= total_len):
+                raise ValueError(
+                    f"Model span {(s, e)} with rest offset {offset_rest} is out of bounds for length {total_len}."
+                )
+            labels[s_abs:e_abs] = torch.tensor(ids_rest[s:e], dtype=torch.long)
+        return labels
+
+    # ---- Trainable summaries -----------------------------------------------------------------
+
+    def summarize_trainables(self) -> Mapping[str, int]:
+        def _count(params: Iterable[nn.Parameter]) -> int:
+            return sum(int(p.numel()) for p in params if p.requires_grad)
+
+        llava_train = _count(self.llava_proj.parameters())
+        ecg_special_train = _count(self.ecg_special_embed.parameters())
+        enc_train = _count(self.enc.parameters())
+        lm_train = _count(self.language_model.parameters())
+
+        total = llava_train + ecg_special_train + enc_train + lm_train
+        return {
+            "llava_proj_trainable": llava_train,
+            "ecg_special_trainable": ecg_special_train,
+            "enc_trainable": enc_train,
+            "lm_trainable": lm_train,
+            "total_trainable": total,
+        }
+
+    # ---- Utility ---------------------------------------------------------------------------
+
+    @staticmethod
+    def ensure_finite(tensor: Tensor, context: str) -> None:
+        if not torch.isfinite(tensor).all():
+            xin_nan = torch.isnan(tensor).any().item()
+            raise RuntimeError(f"Encountered non-finite values in {context} (input_has_nan={bool(xin_nan)}).")

camel_inference/src/camel/ecg_text_packing.py

@@ -0,0 +1,499 @@
+import re
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import List, Dict, Tuple, Optional, Any
+from transformers import PreTrainedTokenizer
+
+from camel.assertions import (
+    assert_ecg_catalog_valid,
+    assert_normalized_role_canonical,
+    assert_turn_parts_structure_valid,
+    assert_turn_content_ends_with_eot,
+    assert_leads_canonical_and_ordered,
+    assert_waveform_shapes_valid,
+)
+from camel.prompt_renderers import turn_wrappers
+
+# NOTE: Local BOS/EOS assertions are implemented at the bottom of this file.
+@dataclass(frozen=True)
+class PromptTokens:
+    start_of_turn: str
+    end_of_turn: str
+    user_role: str
+    model_role: str
+    require_bos: bool = True
+    require_eos: bool = True
+    allow_multiple_eos: bool = False
+
+
+@dataclass(frozen=True)
+class ConversationRules:
+    format_id: str
+    user_role_aliases: Tuple[str, ...] = ("human", "user")
+    model_role_aliases: Tuple[str, ...] = ("gpt", "assistant")
+    strip_image_from_roles: Tuple[str, ...] = ("human",)
+    merge_system_with_first_user: bool = True
+
+
+@dataclass(frozen=True)
+class ECGTokenSchema:
+    global_start: str
+    global_end: str
+    lead_start_template: str
+    lead_end_template: str
+    canonical_leads: Tuple[str, ...]
+
+
+@dataclass(frozen=True)
+class PackingSchema:
+    prompt: PromptTokens
+    conversation: ConversationRules
+    ecg: ECGTokenSchema
+
+
+@dataclass(frozen=True)
+class ECGSpecialTokenCatalog:
+    tokens: Tuple[str, ...]
+    lead_to_indices: Dict[str, Dict[str, int]]
+    lead_to_tokens: Dict[str, Dict[str, str]]
+    token_to_index: Dict[str, int]
+
+
+def _render_lead_template(template: str, lead: str) -> str:
+    return template.format(
+        lead=lead,
+        lead_lower=lead.lower(),
+        lead_upper=lead.upper(),
+    )
+
+_ECG_TOKEN_CACHE: Dict[PackingSchema, ECGSpecialTokenCatalog] = {}
+
+def get_ecg_special_token_catalog(schema: PackingSchema) -> ECGSpecialTokenCatalog:
+    cached = _ECG_TOKEN_CACHE.get(schema)
+    if cached is not None:
+        return cached
+
+    tokens: List[str] = []
+    lead_to_indices: Dict[str, Dict[str, int]] = {}
+    lead_to_tokens: Dict[str, Dict[str, str]] = {}
+
+    tokens.append(schema.ecg.global_start)
+    tokens.append(schema.ecg.global_end)
+
+    for lead in schema.ecg.canonical_leads:
+        start_token = _render_lead_template(schema.ecg.lead_start_template, lead)
+        end_token = _render_lead_template(schema.ecg.lead_end_template, lead)
+        start_idx = len(tokens)
+        tokens.append(start_token)
+        end_idx = len(tokens)
+        tokens.append(end_token)
+        lead_to_indices[lead] = {"start": start_idx, "end": end_idx}
+        lead_to_tokens[lead] = {"start": start_token, "end": end_token}
+
+    catalog = ECGSpecialTokenCatalog(
+        tokens=tuple(tokens),
+        lead_to_indices=lead_to_indices,
+        lead_to_tokens=lead_to_tokens,
+        token_to_index={tok: idx for idx, tok in enumerate(tokens)},
+    )
+
+    # Validate catalog consistency
+    assert_ecg_catalog_valid(catalog, schema)
+
+    _ECG_TOKEN_CACHE[schema] = catalog
+    return catalog
+
+
+# ---- Conversation normalization + validation ----------------------------------------------------
+
+def canonical_leads(schema: PackingSchema) -> List[str]:
+    return list(schema.ecg.canonical_leads)
+
+
+def _strip_image_tag(text: str) -> str:
+    """
+    Remove any <image> placeholder without gluing surrounding words.
+    Replace the token and surrounding whitespace with a single space and
+    normalize repeated spaces.
+    """
+    cleaned = re.sub(r"\s*<image>\s*", " ", text)
+    cleaned = re.sub(r"[ \t]{2,}", " ", cleaned)
+    return cleaned.strip()
+
+
+def _normalize_role(role_value: Any, schema: PackingSchema) -> str:
+    role = (role_value or "").strip()
+    if not role:
+        raise ValueError("Conversation turn is missing a role identifier.")
+    if role.lower() == "system":
+        return "system"
+    if role.lower() == "developer":
+        return "developer"
+    needles = [schema.prompt.user_role] + list(schema.conversation.user_role_aliases)
+    if role.lower() in (needle.lower() for needle in needles):
+        out = schema.prompt.user_role
+        assert_normalized_role_canonical(out, schema)
+        return out
+    needles = [schema.prompt.model_role] + list(schema.conversation.model_role_aliases)
+    if role.lower() in (needle.lower() for needle in needles):
+        out = schema.prompt.model_role
+        assert_normalized_role_canonical(out, schema)
+        return out
+    raise ValueError(f"Unknown conversation role '{role_value}' for schema '{schema.conversation.format_id}'.")
+
+def _normalize_conversation(
+    convo: List[Dict[str, Any]],
+    schema: PackingSchema,
+    system_text: Optional[str],
+    developer_text: Optional[str],
+) -> List[Dict[str, Any]]:
+    if not convo:
+        raise ValueError("Conversation must contain at least one turn.")
+    has_system_turn = False
+    has_developer_turn = False
+    for turn in convo:
+        if not isinstance(turn, dict):
+            continue
+        role_val = turn.get("from")
+        if role_val is None and "role" in turn:
+            role_val = turn.get("role")
+        role_lower = str(role_val or "").strip().lower()
+        if role_lower == "system":
+            has_system_turn = True
+        elif role_lower == "developer":
+            has_developer_turn = True
+    normalized: List[Dict[str, Any]] = []
+    if system_text and system_text.strip() and not has_system_turn:
+        sys = system_text.strip()
+        normalized.append({
+            "role": "system",
+            "content": [{"type": "text", "text": sys}],
+        })
+    if developer_text and developer_text.strip() and not has_developer_turn:
+        dev = developer_text.strip()
+        normalized.append({
+            "role": "developer",
+            "content": [{"type": "text", "text": dev}],
+        })
+    for idx, turn in enumerate(convo):
+        role_val = turn.get("from")
+        if role_val is None:
+            role_val = turn.get("role")
+        role = _normalize_role(role_val, schema)
+        content = turn.get("content")
+        if not isinstance(content, list):
+            raise ValueError(f"Turn {idx} content must be a list of items.")
+        normalized.append({"role": role, "content": content})
+
+    if schema.conversation.merge_system_with_first_user:
+        system_items: List[Dict[str, Any]] = []
+        out: List[Dict[str, Any]] = []
+        first_user_idx: Optional[int] = None
+        for turn in normalized:
+            if turn["role"] == "system":
+                system_items.extend(list(turn["content"]))
+                continue
+            if turn["role"] == "developer":
+                raise ValueError("Developer turn present but merge_system_with_first_user is true.")
+            if first_user_idx is None and turn["role"] == schema.prompt.user_role:
+                first_user_idx = len(out)
+            out.append(turn)
+
+        if system_items:
+            if first_user_idx is None:
+                raise ValueError("System turn present but no user turn to merge into.")
+            user_turn = out[first_user_idx]
+            user_turn["content"] = list(system_items) + list(user_turn["content"])
+
+        if not out:
+            raise ValueError("Conversation must contain at least one non-system turn.")
+        if out[0]["role"] != schema.prompt.user_role:
+            raise ValueError("Conversation must start with a user/human turn.")
+        return out
+
+    out = list(normalized)
+    if not out:
+        raise ValueError("Conversation must contain at least one turn.")
+    system_turns = [t for t in out if t["role"] == "system"]
+    developer_turns = [t for t in out if t["role"] == "developer"]
+    non_preamble_turns = [t for t in out if t["role"] not in ("system", "developer")]
+    if not non_preamble_turns:
+        raise ValueError("Conversation must contain at least one non-system/developer turn.")
+    if non_preamble_turns[0]["role"] != schema.prompt.user_role:
+        raise ValueError("Conversation must start with a user/human turn.")
+    return system_turns + developer_turns + non_preamble_turns
+
+
+def _maybe_strip_content(text: str, canonical_role: str, schema: PackingSchema) -> str:
+    target_roles: set = set()
+    for r in schema.conversation.strip_image_from_roles:
+        try:
+            target_roles.add(_normalize_role(r, schema).lower())
+        except ValueError:
+            target_roles.add(str(r).lower())
+    if canonical_role.lower() in target_roles:
+        return _strip_image_tag(text)
+    return text
+
+# ---- Tokenize & mark assistant spans (exclude control tokens from loss) --------------------------
+
+# ---- Build structured turn parts ---------------------------------------------------------------
+
+def build_structured_turn_parts(
+    *,
+    content: List[Dict[str, Any]],
+    canonical_role: str,
+    schema: PackingSchema,
+    ecg_blocks: List[Dict[str, Any]],
+    sampling_rate: Optional[float],
+    turn_suffix: Optional[str] = None,
+) -> Tuple[str, List[Dict[str, Any]]]:
+    prompt_tokens = schema.prompt
+    catalog = get_ecg_special_token_catalog(schema)
+    parts: List[Dict[str, Any]] = []
+    text_segments: List[str] = []
+
+    def _append_text(txt: str) -> None:
+        if not txt:
+            return
+        if parts and parts[-1].get("kind") == "text":
+            parts[-1]["text"] += txt
+        else:
+            parts.append({"kind": "text", "text": txt})
+        text_segments.append(txt)
+
+    for item in content:
+        if not isinstance(item, dict):
+            raise ValueError("Conversation content items must be dicts.")
+        item_type = item.get("type")
+        if item_type == "text":
+            text_val = item.get("text", "")
+            if not isinstance(text_val, str):
+                raise ValueError("Text content item must have a string 'text' field.")
+            cleaned = _maybe_strip_content(text_val, canonical_role, schema)
+            _append_text(cleaned)
+            continue
+        if item_type == "ecg":
+            waveform_segments = item.get("waveform_segments")
+            if not isinstance(waveform_segments, dict):
+                raise ValueError("ECG content item missing waveform_segments mapping.")
+            item_rate = item.get("sampling_rate")
+            if item_rate is not None and sampling_rate is not None and float(item_rate) != float(sampling_rate):
+                raise ValueError(
+                    f"ECG item sampling_rate {item_rate} does not match sample sampling_rate {sampling_rate}."
+                )
+            lead_names = [str(ld) for ld in waveform_segments.keys()]
+            if not lead_names:
+                raise ValueError("ECG content item has no leads.")
+            segments_per_lead = [int(waveform_segments[ld].shape[0]) for ld in lead_names]
+            assert_leads_canonical_and_ordered(lead_names, schema.ecg.canonical_leads)
+            assert_waveform_shapes_valid(lead_names, segments_per_lead, waveform_segments)
+            block_index = len(ecg_blocks)
+            ecg_blocks.append({
+                "lead_names": lead_names,
+                "segments_per_lead": segments_per_lead,
+                "waveform_segments": OrderedDict((ld, waveform_segments[ld]) for ld in lead_names),
+            })
+
+            parts.append({
+                "kind": "special",
+                "token": schema.ecg.global_start,
+                "token_index": catalog.token_to_index[schema.ecg.global_start],
+                "block_index": block_index,
+            })
+            text_segments.append(schema.ecg.global_start)
+
+            for lead, nseg in zip(lead_names, segments_per_lead):
+                lead_tokens = catalog.lead_to_indices[lead]
+                parts.append({
+                    "kind": "special",
+                    "token": catalog.lead_to_tokens[lead]["start"],
+                    "token_index": lead_tokens["start"],
+                    "lead": lead,
+                    "block_index": block_index,
+                })
+                text_segments.append(catalog.lead_to_tokens[lead]["start"])
+                for sec in range(1, int(nseg) + 1):
+                    parts.append({
+                        "kind": "ecg",
+                        "lead": lead,
+                        "sec": sec,
+                        "block_index": block_index,
+                    })
+                parts.append({
+                    "kind": "special",
+                    "token": catalog.lead_to_tokens[lead]["end"],
+                    "token_index": lead_tokens["end"],
+                    "lead": lead,
+                    "block_index": block_index,
+                })
+                text_segments.append(catalog.lead_to_tokens[lead]["end"])
+
+            parts.append({
+                "kind": "special",
+                "token": schema.ecg.global_end,
+                "token_index": catalog.token_to_index[schema.ecg.global_end],
+                "block_index": block_index,
+            })
+            text_segments.append(schema.ecg.global_end)
+            continue
+        raise ValueError(f"Unknown content item type '{item_type}'.")
+
+    turn_content = "".join(text_segments)
+    if turn_suffix is None:
+        _, suffix = turn_wrappers(schema, canonical_role)
+    else:
+        suffix = turn_suffix
+    turn_text_block = turn_content + suffix
+
+    assert_turn_parts_structure_valid(parts, ecg_blocks, schema, catalog)
+    assert_turn_content_ends_with_eot(turn_text_block, suffix)
+
+    return turn_text_block, parts
+
+
+def build_text_only_turn_parts(
+    *,
+    content: List[Dict[str, Any]],
+    canonical_role: str,
+    schema: PackingSchema,
+    turn_suffix: Optional[str] = None,
+) -> Tuple[str, List[Dict[str, Any]]]:
+    prompt_tokens = schema.prompt
+    parts: List[Dict[str, Any]] = []
+    text_segments: List[str] = []
+    needs_channel_header = (
+        schema.conversation.format_id == "harmony_chat_v1"
+        and canonical_role == schema.prompt.model_role
+    )
+    # Track raw content for diagnostics
+    raw_content_debug: List[Dict[str, Any]] = []
+
+    def _append_text(txt: str) -> None:
+        if not txt:
+            return
+        if parts and parts[-1].get("kind") == "text":
+            parts[-1]["text"] += txt
+        else:
+            parts.append({"kind": "text", "text": txt})
+        text_segments.append(txt)
+
+    for item in content:
+        if not isinstance(item, dict):
+            raise ValueError("Conversation content items must be dicts.")
+        item_type = item.get("type")
+        if item_type != "text":
+            raise ValueError("Model turns cannot contain ECG content items.")
+        text_val = item.get("text", "")
+        if not isinstance(text_val, str):
+            raise ValueError("Text content item must have a string 'text' field.")
+        if needs_channel_header and text_val.lstrip().startswith("<|channel|>"):
+            raise ValueError("Assistant content must not include harmony channel headers.")
+        cleaned = _maybe_strip_content(text_val, canonical_role, schema)
+        raw_content_debug.append({
+            "raw_text": repr(text_val[:200]) + ("..." if len(text_val) > 200 else ""),
+            "raw_len": len(text_val),
+            "cleaned_text": repr(cleaned[:200]) + ("..." if len(cleaned) > 200 else ""),
+            "cleaned_len": len(cleaned),
+        })
+        _append_text(cleaned)
+
+    turn_content = "".join(text_segments)
+    if not turn_content:
+        # Build detailed diagnostic message
+        diag_lines = [
+            "Model turn is empty after preprocessing.",
+            f"  role: {canonical_role}",
+            f"  num_content_items: {len(content)}",
+        ]
+        for i, dbg in enumerate(raw_content_debug):
+            diag_lines.append(f"  item[{i}]: raw_len={dbg['raw_len']}, cleaned_len={dbg['cleaned_len']}")
+            diag_lines.append(f"    raw: {dbg['raw_text']}")
+            diag_lines.append(f"    cleaned: {dbg['cleaned_text']}")
+        raise ValueError("\n".join(diag_lines))
+    if needs_channel_header:
+        channel_header = "<|channel|>final<|message|>"
+        parts.insert(0, {"kind": "text", "text": channel_header})
+        text_segments.insert(0, channel_header)
+        turn_content = "".join(text_segments)
+    if turn_suffix is None:
+        _, suffix = turn_wrappers(schema, canonical_role)
+    else:
+        suffix = turn_suffix
+    turn_text_block = turn_content + suffix
+    assert_turn_content_ends_with_eot(turn_text_block, suffix)
+    return turn_text_block, parts
+
+
+def annotate_turn_parts_with_ids(
+    turn_parts: List[List[Dict[str, Any]]],
+    tokenizer: PreTrainedTokenizer,
+) -> List[List[Dict[str, Any]]]:
+    """Attach token ids to text parts so the trainer can skip per-step tokenization."""
+    for parts in turn_parts:
+        for part in parts:
+            if part.get("kind") == "text":
+                txt = part.get("text", "")
+                part["ids"] = tokenizer.encode(txt, add_special_tokens=False) if txt else []
+    return turn_parts
+
+def assert_single_bos_eos(
+    text_ids: List[int],
+    tok: PreTrainedTokenizer,
+    *,
+    require_bos_at_start: bool,
+    require_single_terminal_eos: bool,
+    allow_multiple_eos: bool = False,
+) -> None:
+    """Validate BOS/EOS placement under various schema-specific policies."""
+    if require_bos_at_start:
+        if tok.bos_token_id is None:
+            raise AssertionError("BOS token required but tokenizer has none")
+        bos_pos = [i for i, t in enumerate(text_ids) if t == tok.bos_token_id]
+        if len(bos_pos) != 1 or bos_pos[0] != 0:
+            raise AssertionError(f"BOS placement invalid: positions={bos_pos}")
+    else:
+        if tok.bos_token_id is not None:
+            bos_pos = [i for i, t in enumerate(text_ids) if t == tok.bos_token_id]
+            if len(bos_pos) > 1 or (bos_pos and bos_pos[0] != 0):
+                raise AssertionError(f"Unexpected BOS placement: positions={bos_pos}")
+
+    # EOS checks
+    if tok.eos_token_id is not None:
+        eos_pos = [i for i, t in enumerate(text_ids) if t == tok.eos_token_id]
+        if allow_multiple_eos:
+            if require_single_terminal_eos:
+                # Allow multiple EOS (e.g., ChatML-style per turn), but require the last EOS at sequence end.
+                if len(eos_pos) == 0 or eos_pos[-1] != (len(text_ids) - 1):
+                    raise AssertionError(f"EOS bad: positions={eos_pos}")
+            else:
+                # Allow any count anywhere; no terminal EOS required (matches Qwen3 training practice).
+                pass
+        else:
+            if require_single_terminal_eos:
+                # Require exactly one EOS, and it must be terminal.
+                if len(eos_pos) != 1 or eos_pos[0] != (len(text_ids) - 1):
+                    raise AssertionError(f"EOS bad: positions={eos_pos}")
+            else:
+                # Allow 0 or 1, but if present it must be terminal.
+                if len(eos_pos) > 1 or (eos_pos and eos_pos[0] != (len(text_ids) - 1)):
+                    raise AssertionError(f"EOS bad: positions={eos_pos}")
+
+
+def assert_struct_bos_eos(
+    token_struct: Dict[str, List[int]],
+    tok: PreTrainedTokenizer,
+    *,
+    require_bos_at_start: bool,
+    require_single_terminal_eos: bool,
+    allow_multiple_eos: bool = False,
+) -> None:
+    """Wrapper that validates concatenated ids per schema policy."""
+    text_ids = token_struct["text_ids"]
+    assert_single_bos_eos(
+        text_ids,
+        tok,
+        require_bos_at_start=require_bos_at_start,
+        require_single_terminal_eos=require_single_terminal_eos,
+        allow_multiple_eos=allow_multiple_eos,
+    )

camel_inference/src/camel/inference.py

@@ -0,0 +1,846 @@
+#!/usr/bin/env python3
+# inference.py — ECGText inference (model loading + generation helpers)
+import os
+import re
+from dataclasses import dataclass
+from typing import Dict, List, Tuple, Optional, Any
+from collections import OrderedDict
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers import AutoModelForCausalLM
+from peft import LoraConfig
+
+# Local imports
+from camel.model_introspect import resolve_hidden_size as _resolve_hidden_size
+from camel.model_registry import load_registry
+from camel.training_setup import initialize_tokenizer, build_packing_schema, register_ecg_special_tokens
+from camel.model_init import build_wrapper, attach_lora, build_conv_encoder
+from camel.ecg_text_packing import (
+    _normalize_conversation,
+    annotate_turn_parts_with_ids,
+    build_structured_turn_parts,
+    build_text_only_turn_parts,
+    get_ecg_special_token_catalog,
+)
+from camel.prompt_renderers import render_prompt_and_spans, turn_wrappers, assistant_generation_prefix
+from camel.ecg_attention_masks import (
+    ECGBlockLayout,
+    ECGSequenceLayout,
+    MaskBuildResult,
+    ECGMaskStrategy,
+    get_mask_strategy,
+)
+from camel.assertions import (
+    assert_ecg_blocks_consistent,
+    assert_ecg_part_bounds,
+    assert_layout_specials_complete,
+    assert_prefix_matches_segments,
+    assert_prefix_split_complete,
+)
+from camel.checkpoint_utils import (
+    load_llava_and_lora,
+    update_wrapper_language_model,
+    extract_lora_config_from_checkpoints,
+    peek_projector_name,
+)
+
+# ------------------------------
+# Device & conv builder
+# ------------------------------
+
+def _device(device=None) -> torch.device:
+    if device:
+        return device
+    if torch.cuda.is_available():
+        local_rank = int(os.environ.get("LOCAL_RANK", 0))
+        return torch.device(f"cuda:{local_rank}")
+    return torch.device("cpu")
+
+
+_HARMONY_CHANNEL_RE = re.compile(r"<\|channel\|>(.*?)<\|message\|>", re.DOTALL)
+_HARMONY_DELIM_RE = re.compile(r"<\|end\|>|<\|return\|>|<\|call\|>|<\|start\|>")
+
+
+def _extract_harmony_messages(text: str) -> List[Tuple[str, str]]:
+    matches = list(_HARMONY_CHANNEL_RE.finditer(text))
+    if not matches:
+        raise ValueError("No harmony channel headers found in model output.")
+    out: List[Tuple[str, str]] = []
+    for match in matches:
+        channel_raw = match.group(1).strip()
+        channel = channel_raw.split()[0] if channel_raw else ""
+        if not channel:
+            raise ValueError("Harmony channel header is empty.")
+        start = match.end()
+        end_match = _HARMONY_DELIM_RE.search(text, start)
+        end = end_match.start() if end_match else len(text)
+        out.append((channel, text[start:end]))
+    return out
+
+
+def _checkpoint_has_conv(ckpt_path: Optional[str]) -> bool:
+    if not ckpt_path:
+        return False
+    payload = torch.load(ckpt_path, map_location="cpu")
+    if not isinstance(payload, dict):
+        raise RuntimeError(f"Checkpoint {ckpt_path} must be a dict to inspect conv metadata.")
+    return isinstance(payload.get("conv"), dict)
+
+
+# ------------------------------
+# Prompt building & stopping
+# ------------------------------
+
+@dataclass
+class PromptContext:
+    """Container describing the prepared prompt state for autoregressive generation."""
+
+    inputs_embeds: torch.Tensor
+    layout: ECGSequenceLayout
+    prompt_preview: str
+    stop_ids: List[int]
+    input_embedder: nn.Embedding
+    mask_strategy: ECGMaskStrategy
+    mask_result: MaskBuildResult
+
+
+def _sanitize_segments(tensor: torch.Tensor) -> torch.Tensor:
+    """Detach → float32 → replace NaN/Inf so downstream encoders stay numerically stable."""
+    out = tensor.detach().cpu().to(dtype=torch.float32)
+    return torch.nan_to_num(out, nan=0.0, posinf=0.0, neginf=0.0)
+
+
+def _sample_next_token(
+    logits: torch.Tensor,
+    *,
+    temperature: float,
+    top_k: Optional[int],
+    top_p: float,
+    min_p: float,
+) -> torch.Tensor:
+    """
+    Draw the next token id given final-step logits and sampling parameters.
+
+    Uses greedy decoding when temperature <= 0; otherwise applies temperature
+    scaling, optional nucleus sampling, and multinomial sampling.
+    """
+    if logits.ndim != 1:
+        raise ValueError(f"Expected 1D logits, got shape {tuple(logits.shape)}")
+
+    if temperature <= 0.0:
+        return torch.argmax(logits, dim=-1)
+
+    scaled = logits / max(temperature, 1e-5)
+    probs = torch.softmax(scaled, dim=-1)
+
+    sorted_probs, sorted_indices = torch.sort(probs, descending=True)
+    keep_mask = torch.ones_like(sorted_probs, dtype=torch.bool)
+
+    if top_k is not None and top_k > 0:
+        top_k = min(int(top_k), sorted_probs.numel())
+        top_k_mask = torch.zeros_like(sorted_probs, dtype=torch.bool)
+        top_k_mask[:top_k] = True
+        keep_mask &= top_k_mask
+
+    if 0.0 < top_p < 1.0:
+        cumulative = torch.cumsum(sorted_probs, dim=-1)
+        cutoff_mask = (cumulative - sorted_probs) < top_p
+        cutoff_mask[0] = True  # always keep the highest-prob token
+        keep_mask &= cutoff_mask
+
+    if min_p is not None and min_p > 0.0:
+        keep_mask &= sorted_probs >= float(min_p)
+
+    filtered_probs = sorted_probs[keep_mask]
+    filtered_indices = sorted_indices[keep_mask]
+    if filtered_probs.numel() == 0:
+        filtered_probs = sorted_probs[:1]
+        filtered_indices = sorted_indices[:1]
+
+    prob_sum = filtered_probs.sum()
+    if not torch.isfinite(prob_sum) or prob_sum <= 0:
+        return sorted_indices[0]
+    normalized = filtered_probs / prob_sum
+    next_idx = torch.multinomial(normalized, num_samples=1, replacement=False)
+    return filtered_indices[next_idx].squeeze(0)
+
+
+class KardiaLM:
+    """High-level chat interface around an ECG language model."""
+
+    def __init__(
+        self,
+        *,
+        model_registry_path: Optional[str],
+        model_config_name: str,
+        hf_model_id_override: Optional[str],
+        adapter_ckpt: str,
+        conv_ckpt: Optional[str] = None,
+        no_lora: bool = False,
+        use_dora: bool = False,
+        default_max_new_tokens: int = 1000,
+        default_temperature: float = 1.0,
+        default_top_k: Optional[int] = 64,
+        default_top_p: float = 0.95,
+        default_min_p: float = 0.0,
+        mask_strategy: Optional[str] = None,
+        device: Optional[torch.device] = None,
+    ) -> None:
+        registry = load_registry(registry_path=model_registry_path)
+        model_cfg = registry.get(model_config_name)
+
+        self.model_cfg = model_cfg
+        self.hf_model_id = hf_model_id_override or model_cfg.hf_id
+        self.packing_schema = build_packing_schema(self.hf_model_id)
+        self.tokenizer_cfg = model_cfg.tokenizer_config()
+        self.arch_cfg = model_cfg.architecture_config()
+        self.system_text = None
+        self.developer_text = None
+        if self.packing_schema.conversation.format_id == "harmony_chat_v1":
+            self.system_text = model_cfg.required_prompt_text("system_prompt")
+            self.developer_text = model_cfg.required_prompt_text("developer_prompt")
+            if not self.system_text.strip():
+                raise RuntimeError("System prompt text for harmony format must be non-empty.")
+            if not self.developer_text.strip():
+                raise RuntimeError("Developer prompt text for harmony format must be non-empty.")
+
+        self.device = _device(device)
+        self.dtype = torch.bfloat16
+        self.mask_strategy: ECGMaskStrategy = get_mask_strategy(mask_strategy)
+        self.expect_dora = bool(use_dora)
+
+        tok = initialize_tokenizer(
+            self.hf_model_id,
+            trust_remote_code=True,
+            use_fast=self.tokenizer_cfg.use_fast,
+            add_prefix_space=self.tokenizer_cfg.add_prefix_space,
+        )
+        self.tokenizer = tok
+
+        catalog = get_ecg_special_token_catalog(self.packing_schema)
+        self.ecg_special_token_id_map = register_ecg_special_tokens(tok, catalog)
+
+        pad_strategy = self.tokenizer_cfg.pad_token_strategy.lower()
+        if pad_strategy == "eos":
+            if tok.eos_token is None:
+                raise RuntimeError(
+                    f"Tokenizer for model '{model_cfg.name}' lacks an EOS token required for pad_token_strategy='eos'."
+                )
+            tok.pad_token = tok.eos_token
+        elif pad_strategy not in ("existing", "keep"):
+            raise RuntimeError(f"Unsupported pad_token_strategy '{self.tokenizer_cfg.pad_token_strategy}'.")
+
+        if self.tokenizer_cfg.require_bos and tok.bos_token is None:
+            raise RuntimeError(f"Tokenizer for model '{model_cfg.name}' is missing a BOS token.")
+        if self.tokenizer_cfg.require_eos and tok.eos_token is None:
+            raise RuntimeError(f"Tokenizer for model '{model_cfg.name}' is missing an EOS token.")
+
+        attn_impl = self.arch_cfg.attn_implementation or "flash_attention_2"
+        try:
+            model = AutoModelForCausalLM.from_pretrained(
+                self.hf_model_id,
+                torch_dtype=self.dtype,
+                trust_remote_code=True,
+                attn_implementation=attn_impl,
+                device_map=None,
+            ).to(self.device)
+        except Exception:
+            model = AutoModelForCausalLM.from_pretrained(
+                self.hf_model_id,
+                torch_dtype=self.dtype,
+                trust_remote_code=True,
+                attn_implementation="eager",
+                device_map=None,
+            ).to(self.device)
+        if model.get_input_embeddings().weight.shape[0] != len(tok):
+            model.resize_token_embeddings(len(tok))
+        for p in model.parameters():
+            p.requires_grad = False
+        model.eval()
+        if hasattr(model, "gradient_checkpointing_disable"):
+            model.gradient_checkpointing_disable()
+        if hasattr(model.config, "use_cache"):
+            model.config.use_cache = True
+        self.model = model
+
+        adapter_ckpt_path = os.path.expanduser(adapter_ckpt)
+        adapter_has_conv = _checkpoint_has_conv(adapter_ckpt_path)
+        if not adapter_has_conv and not conv_ckpt:
+            raise RuntimeError(
+                "Adapter checkpoint lacks conv weights; supply --conv_ckpt to match training."
+            )
+        lora_cfg_dict = extract_lora_config_from_checkpoints(adapter_ckpt_path, None)
+        active_lora_cfg: Optional[LoraConfig] = None
+        if self.expect_dora and no_lora:
+            raise RuntimeError("--use-dora cannot be combined with --no-lora since no adapters would be loaded.")
+        if lora_cfg_dict and not no_lora:
+            cfg_use_dora = bool(lora_cfg_dict.get("use_dora", False))
+            if cfg_use_dora and not self.expect_dora:
+                raise RuntimeError(
+                    "Checkpoint adapters were trained with DoRA; re-run inference with --use-dora to load them."
+                )
+            if self.expect_dora and not cfg_use_dora:
+                raise RuntimeError(
+                    "Checkpoint adapters were trained without DoRA; omit --use-dora or use a checkpoint with DoRA."
+                )
+            model, active_lora_cfg = attach_lora(model, lora_cfg_dict, self.device)
+            model.eval()
+        elif no_lora and lora_cfg_dict:
+            print("[LoRA] --no-lora set; skipping LoRA adapters from checkpoint.", flush=True)
+        elif self.expect_dora:
+            raise RuntimeError("--use-dora was provided, but no LoRA/DoRA adapters were found in the checkpoint.")
+
+        conv = build_conv_encoder(
+            conv_ckpt_path=None if adapter_has_conv else conv_ckpt,
+            device=self.device,
+            unfreeze=False,
+        )
+        conv.eval()
+        for p in conv.parameters():
+            p.requires_grad = False
+        self.conv_encoder = conv
+
+        hidden_size = _resolve_hidden_size(model, self.arch_cfg.hidden_size_attrs)
+        wrapper_cls = model_cfg.resolve_wrapper_class()
+        enc_out_dim = self.arch_cfg.conv_out_dim if getattr(self.arch_cfg, "conv_out_dim", None) is not None else 64
+        projector_name = peek_projector_name(adapter_ckpt_path) or "linear"
+        wrapper = build_wrapper(
+            wrapper_cls=wrapper_cls,
+            language_model=model,
+            conv_encoder=conv,
+            hidden_size=hidden_size,
+            num_ecg_special_tokens=len(catalog.tokens),
+            dtype=self.dtype,
+            enc_out_dim=int(enc_out_dim),
+            freeze_encoder=True,
+            inference=True,
+            projector_name=projector_name,
+        )
+        self.projector_name = projector_name
+        self.wrapper = wrapper
+
+        _extra_payload, model, inferred_lora_cfg = load_llava_and_lora(
+            wrapper,
+            model,
+            adapter_ckpt_path,
+            expect_lora=(active_lora_cfg is not None),
+            load_lora=not no_lora,
+        )
+        update_wrapper_language_model(wrapper, model)
+        if active_lora_cfg is None and inferred_lora_cfg is not None and not no_lora:
+            active_lora_cfg = inferred_lora_cfg
+        model.eval()
+        for p in model.parameters():
+            p.requires_grad = False
+
+        inp_emb = model.get_input_embeddings().weight
+        inp_dev = inp_emb.device
+        target_dtype = inp_emb.dtype
+        wrapper.llava_proj.to(device=inp_dev, dtype=torch.float32)
+        wrapper.enc.to(device=inp_dev, dtype=torch.float32)
+        wrapper.ecg_special_embed.to(device=inp_dev, dtype=target_dtype)
+        llava_param = next(wrapper.llava_proj.parameters(), None)
+        if llava_param is None:
+            raise AssertionError("llava_proj unexpectedly has no parameters.")
+        if llava_param.device != inp_dev:
+            raise AssertionError(f"llava_proj on {llava_param.device}, expected {inp_dev}")
+        if llava_param.dtype != torch.float32:
+            raise AssertionError(f"llava_proj dtype {llava_param.dtype}, expected torch.float32")
+        conv_param = next(wrapper.enc.parameters(), None)
+        if conv_param is None:
+            raise AssertionError("Convolutional encoder unexpectedly has no parameters.")
+        if conv_param.device != inp_dev:
+            raise AssertionError(f"Conv encoder on {conv_param.device}, expected {inp_dev}")
+        if conv_param.dtype != torch.float32:
+            raise AssertionError(f"Conv encoder dtype {conv_param.dtype}, expected torch.float32")
+        assert next(wrapper.ecg_special_embed.parameters()).device == inp_dev, (
+            f"ecg_special_embed on {next(wrapper.ecg_special_embed.parameters()).device}, expected {inp_dev}"
+        )
+        try:
+            wrapper.language_model.eval()
+        except Exception:
+            pass
+
+        self.default_max_new_tokens = int(default_max_new_tokens)
+        self.default_temperature = float(default_temperature)
+        self.default_top_k = int(default_top_k) if default_top_k is not None else None
+        self.default_top_p = float(default_top_p)
+        self.default_min_p = float(default_min_p)
+
+    def chat(
+        self,
+        *,
+        conversation: List[Dict[str, Any]],
+        max_new_tokens: Optional[int] = None,
+        temperature: Optional[float] = None,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+        min_p: Optional[float] = None,
+        harmony_output: Optional[str] = None,
+    ) -> Tuple[str, str]:
+        """Generate a response for a structured multi-turn conversation."""
+        context = self._prepare_prompt_context(conversation=conversation)
+
+        max_new_tokens = int(max_new_tokens if max_new_tokens is not None else self.default_max_new_tokens)
+        temperature = float(temperature if temperature is not None else self.default_temperature)
+        resolved_top_k = int(top_k) if top_k is not None else (self.default_top_k if self.default_top_k is not None else None)
+        top_p = float(top_p if top_p is not None else self.default_top_p)
+        min_p = float(min_p if min_p is not None else self.default_min_p)
+
+        token_ids = self._autoregressive_generate(
+            context=context,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_k=resolved_top_k,
+            top_p=top_p,
+            min_p=min_p,
+        )
+        text = self.tokenizer.decode(token_ids, skip_special_tokens=False)
+        if self.packing_schema.conversation.format_id == "harmony_chat_v1":
+            mode = harmony_output if harmony_output is not None else "all"
+            if mode != "raw":
+                messages = _extract_harmony_messages(text)
+                if mode == "all":
+                    text = "\n".join(msg for _, msg in messages)
+                elif mode == "final":
+                    finals = [msg for channel, msg in messages if channel == "final"]
+                    if not finals:
+                        raise ValueError("No final channel output found in harmony response.")
+                    text = finals[-1]
+                else:
+                    raise ValueError(f"Unknown harmony_output '{mode}'.")
+        return text, context.prompt_preview
+
+    def _to_waveform_tensor(self, value: Any) -> torch.Tensor:
+        if isinstance(value, torch.Tensor):
+            tensor = value.detach().cpu()
+        elif isinstance(value, np.ndarray):
+            tensor = torch.from_numpy(np.asarray(value))
+        else:
+            tensor = torch.tensor(value, dtype=torch.float32)
+        tensor = tensor.to(dtype=torch.float32)
+        if tensor.ndim == 1:
+            if tensor.numel() != 256:
+                raise ValueError("Expected a 256-sample vector for a single lead second.")
+            tensor = tensor.view(1, 256)
+        elif tensor.ndim == 2:
+            if tensor.size(-1) != 256:
+                raise ValueError("Waveform segments must have length 256 along the last dimension.")
+        else:
+            raise ValueError("Waveform tensor must be rank 1 or 2 with 256-sample segments.")
+        return tensor.contiguous()
+
+    def _prepare_prompt_context(
+        self,
+        *,
+        conversation: List[Dict[str, Any]],
+    ) -> PromptContext:
+        tok = self.tokenizer
+        wrapper = self.wrapper
+        packing_schema = self.packing_schema
+        device = self.device
+
+        prompt_tokens = packing_schema.prompt
+        if not isinstance(conversation, list) or not conversation:
+            raise ValueError("conversation must be a non-empty list of turns.")
+
+        conv_input: List[Dict[str, Any]] = []
+        for turn in conversation:
+            if not isinstance(turn, dict):
+                raise ValueError("Conversation turns must be dicts.")
+            if "from" not in turn and "role" in turn:
+                turn = dict(turn)
+                turn["from"] = turn.get("role")
+            conv_input.append(turn)
+
+        turns = _normalize_conversation(conv_input, packing_schema, self.system_text, self.developer_text)
+        if turns[-1]["role"] != prompt_tokens.user_role:
+            raise ValueError("Conversation must end with a user turn to generate.")
+
+        def _sanitize_content(content: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+            sanitized: List[Dict[str, Any]] = []
+            for item in content:
+                if not isinstance(item, dict):
+                    raise ValueError("Conversation content items must be dicts.")
+                item_type = item.get("type")
+                if item_type == "ecg":
+                    waveform = item.get("waveform_segments")
+                    if not isinstance(waveform, dict):
+                        raise ValueError("ECG content item missing waveform_segments mapping.")
+                    wf_out: "OrderedDict[str, torch.Tensor]" = OrderedDict()
+                    for ld, value in waveform.items():
+                        wf_out[str(ld)] = _sanitize_segments(self._to_waveform_tensor(value))
+                    new_item = dict(item)
+                    new_item["waveform_segments"] = wf_out
+                    sanitized.append(new_item)
+                    continue
+                if item_type == "text":
+                    text_val = item.get("text")
+                    if not isinstance(text_val, str):
+                        raise ValueError("Text content item must have a string 'text' field.")
+                    if "<image>" in text_val:
+                        raise ValueError("Conversation text must not contain <image> in inference mode.")
+                sanitized.append(item)
+            return sanitized
+
+        ecg_blocks: List[Dict[str, Any]] = []
+        turn_parts: List[List[Dict[str, Any]]] = []
+        token_turns: List[Dict[str, str]] = []
+
+        for turn in turns:
+            role = turn["role"]
+            content = _sanitize_content(turn["content"])
+            if role == prompt_tokens.model_role:
+                turn_text_block, content_parts = build_text_only_turn_parts(
+                    content=content,
+                    canonical_role=role,
+                    schema=packing_schema,
+                )
+            else:
+                turn_text_block, content_parts = build_structured_turn_parts(
+                    content=content,
+                    canonical_role=role,
+                    schema=packing_schema,
+                    ecg_blocks=ecg_blocks,
+                    sampling_rate=None,
+                )
+            prefix, suffix = turn_wrappers(packing_schema, role)
+            parts = [{"kind": "text", "text": prefix}]
+            parts.extend(content_parts)
+            parts.append({"kind": "text", "text": suffix})
+            turn_parts.append(parts)
+            token_turns.append({"role": role, "text_block": turn_text_block})
+
+        if not ecg_blocks:
+            raise ValueError("No ECG blocks found in conversation.")
+
+        turn_parts = annotate_turn_parts_with_ids(turn_parts, tok)
+        assert_ecg_blocks_consistent(turn_parts=turn_parts, ecg_blocks=ecg_blocks)
+
+        token_struct = render_prompt_and_spans(tok, token_turns, schema=packing_schema)
+        text_ids = list(token_struct["text_ids"])
+        if (
+            prompt_tokens.require_eos
+            and tok.eos_token_id is not None
+            and text_ids
+            and text_ids[-1] == tok.eos_token_id
+        ):
+            text_ids = text_ids[:-1]
+        text_preview = token_struct.get("text_preview", "")
+
+        model_prefix = assistant_generation_prefix(packing_schema)
+        ids_model_prefix = tok.encode(model_prefix, add_special_tokens=False)
+        text_ids.extend(ids_model_prefix)
+        prompt_preview = text_preview + model_prefix
+
+        model_prefix_parts = [{"kind": "text", "text": model_prefix, "ids": ids_model_prefix}]
+        all_parts = list(turn_parts) + [model_prefix_parts]
+
+        flat_blocks = [blk["waveform_segments"] for blk in ecg_blocks]
+        lead_orders = [blk["lead_names"] for blk in ecg_blocks]
+        prefix_all, prefix_lens = wrapper.ecg_prefix_batch(
+            flat_blocks,
+            device=device,
+            lead_orders=lead_orders,
+        )
+        prefixes: List[torch.Tensor] = []
+        offset = 0
+        for n in prefix_lens:
+            prefixes.append(prefix_all[offset:offset + int(n)])
+            offset += int(n)
+        assert_prefix_split_complete(offset=offset, total_prefix_rows=int(prefix_all.size(0)))
+
+        block_layouts: List[ECGBlockLayout] = []
+        lead_offsets: List[Dict[str, int]] = []
+        lead_special_counts: List[Dict[str, int]] = []
+        for blk_idx, blk in enumerate(ecg_blocks):
+            lead_names = [str(ld) for ld in blk.get("lead_names", [])]
+            segs_per_lead = [int(n) for n in blk.get("segments_per_lead", [])]
+            prefix_rows = prefixes[blk_idx].size(0) if blk_idx < len(prefixes) else 0
+            assert_prefix_matches_segments(
+                prefix_rows=prefix_rows,
+                segments_per_lead=segs_per_lead,
+                lead_names=lead_names,
+                sample_index=0,
+                block_index=blk_idx,
+            )
+            lead_to_offset: Dict[str, int] = {}
+            c = 0
+            for ld, nseg in zip(lead_names, segs_per_lead):
+                lead_to_offset[ld] = c
+                c += int(nseg)
+            lead_offsets.append(lead_to_offset)
+            block_layouts.append(ECGBlockLayout(
+                start_idx=None,
+                end_idx_exclusive=None,
+                global_start_idx=None,
+                global_end_idx=None,
+                lead_start_idx={},
+                lead_end_idx={},
+                signal_pos_by_lead={ld: [None] * int(nseg) for ld, nseg in zip(lead_names, segs_per_lead)},
+                time_to_signal_idxs={},
+                declared_segments_per_lead={ld: int(nseg) for ld, nseg in zip(lead_names, segs_per_lead)},
+            ))
+            lead_special_counts.append({})
+
+        special_indices: List[int] = [
+            int(part["token_index"])
+            for turn in all_parts
+            for part in turn
+            if part.get("kind") == "special"
+        ]
+        if special_indices:
+            special_idx_tensor = torch.tensor(special_indices, dtype=torch.long, device=device)
+            special_embeds = wrapper.ecg_special_tokens_to_embeds(special_idx_tensor, device=device)
+        else:
+            special_embeds = torch.empty((0, wrapper.hidden_size), dtype=wrapper.dtype, device=device)
+
+        input_embedder = wrapper.language_model.get_input_embeddings()
+        if text_ids:
+            E_text_all = wrapper.tokens_to_embeds(input_embedder, text_ids, device=device)
+        else:
+            E_text_all = torch.empty((0, wrapper.hidden_size), dtype=wrapper.dtype, device=device)
+
+        text_cursor = 0
+        empty_text = E_text_all[:0]
+        chunks: List[torch.Tensor] = []
+        layout = ECGSequenceLayout(seq_len=0, text_idxs=[], blocks=block_layouts)
+
+        def _take_text(count: int) -> torch.Tensor:
+            nonlocal text_cursor
+            if count <= 0:
+                return empty_text
+            end = text_cursor + count
+            if end > E_text_all.size(0):
+                raise RuntimeError("Text embedding cursor exceeded available embeddings")
+            out = E_text_all[text_cursor:end]
+            text_cursor = end
+            return out
+
+        def _record_text(count: int, cursor: int) -> None:
+            for i in range(count):
+                layout.text_idxs.append(cursor + i)
+
+        cursor = 0
+        special_cursor = 0
+
+        if (
+            text_ids
+            and prompt_tokens.require_bos
+            and tok.bos_token_id is not None
+            and text_ids[0] == tok.bos_token_id
+        ):
+            E_bos = _take_text(1)
+            chunks.append(E_bos)
+            _record_text(1, cursor)
+            cursor += 1
+
+        for turn in all_parts:
+            for part in turn:
+                kind = part.get("kind")
+                if kind == "text":
+                    ids_chunk = part.get("ids")
+                    if ids_chunk is None:
+                        txt = part.get("text", "")
+                        ids_chunk = tok.encode(txt, add_special_tokens=False) if txt else []
+                    if ids_chunk:
+                        if ids_chunk != text_ids[text_cursor:text_cursor + len(ids_chunk)]:
+                            raise RuntimeError("Special token id does not match text_ids cursor.")
+                        E_chunk = _take_text(len(ids_chunk))
+                        chunks.append(E_chunk)
+                        _record_text(len(ids_chunk), cursor)
+                        cursor += len(ids_chunk)
+                    continue
+                if kind == "special":
+                    if special_cursor >= special_embeds.size(0):
+                        raise RuntimeError("Special-token cursor exceeded embeddings.")
+                    tok_idx = int(part.get("token_index", -1))
+                    expected_id = self.ecg_special_token_id_map.get(tok_idx)
+                    if expected_id is None:
+                        raise RuntimeError(f"Unknown ECG special token index {tok_idx}.")
+                    if text_cursor >= len(text_ids):
+                        raise RuntimeError("Text cursor exceeded available text ids.")
+                    if text_ids[text_cursor] != expected_id:
+                        raise RuntimeError("Special token id does not match text_ids cursor.")
+                    _take_text(1)
+                    chunks.append(special_embeds[special_cursor:special_cursor + 1])
+                    _record_text(1, cursor)
+
+                    block_index = int(part.get("block_index", -1))
+                    if block_index < 0 or block_index >= len(block_layouts):
+                        raise RuntimeError("ECG part references unknown block_index.")
+                    block_layout = block_layouts[block_index]
+                    lead_name = part.get("lead")
+                    if lead_name:
+                        cnt = lead_special_counts[block_index].get(lead_name, 0)
+                        if cnt == 0:
+                            block_layout.lead_start_idx[lead_name] = cursor
+                        else:
+                            block_layout.lead_end_idx[lead_name] = cursor
+                        lead_special_counts[block_index][lead_name] = cnt + 1
+                    else:
+                        if block_layout.global_start_idx is None:
+                            block_layout.global_start_idx = cursor
+                            block_layout.start_idx = cursor
+                        else:
+                            block_layout.global_end_idx = cursor
+                            block_layout.end_idx_exclusive = cursor + 1
+
+                    cursor += 1
+                    special_cursor += 1
+                    continue
+                if kind == "ecg":
+                    block_index = int(part.get("block_index", -1))
+                    if block_index < 0 or block_index >= len(block_layouts):
+                        raise RuntimeError("ECG part references unknown block_index.")
+                    ld = part["lead"]
+                    sec = int(part["sec"])
+                    lead_to_offset = lead_offsets[block_index]
+                    block_layout = block_layouts[block_index]
+                    prefix_all = prefixes[block_index]
+                    assert_ecg_part_bounds(
+                        lead=ld,
+                        sec=sec,
+                        lead_to_offset=lead_to_offset,
+                        declared_segments=block_layout.declared_segments_per_lead,
+                        total_prefix_rows=prefix_all.size(0),
+                        sample_index=0,
+                        block_index=block_index,
+                    )
+                    base = lead_to_offset[ld]
+                    row_idx = base + (sec - 1)
+                    chunks.append(prefix_all[row_idx:row_idx + 1])
+                    sig_list = block_layout.signal_pos_by_lead[ld]
+                    if sec - 1 >= len(sig_list):
+                        raise RuntimeError("ECG segment index exceeds declared segments_per_lead")
+                    sig_list[sec - 1] = cursor
+                    block_layout.time_to_signal_idxs.setdefault(sec, []).append(cursor)
+                    cursor += 1
+                    continue
+                raise RuntimeError(f"Unknown turn part kind '{kind}'.")
+
+        remaining = len(text_ids) - text_cursor
+        if remaining > 0:
+            E_tail = _take_text(remaining)
+            chunks.append(E_tail)
+            _record_text(remaining, cursor)
+            cursor += remaining
+
+        if special_cursor != special_embeds.size(0):
+            raise RuntimeError("Did not consume all special-token embeddings for prompt")
+        if text_cursor != E_text_all.size(0):
+            raise RuntimeError("Text embedding cursor did not consume all embeddings")
+
+        inputs_embeds = torch.cat(chunks, dim=0)
+        layout.seq_len = inputs_embeds.size(0)
+
+        for blk_idx, blk_layout in enumerate(block_layouts):
+            for ld, expected in blk_layout.declared_segments_per_lead.items():
+                slots = blk_layout.signal_pos_by_lead[ld]
+                if any(pos is None for pos in slots):
+                    raise RuntimeError(f"Lead {ld} missing ECG slots; expected {expected}.")
+                blk_layout.signal_pos_by_lead[ld] = [int(pos) for pos in slots]
+            if blk_layout.global_start_idx is None or blk_layout.global_end_idx is None:
+                raise RuntimeError("ECG block missing global start/end specials.")
+            if blk_layout.end_idx_exclusive is None:
+                blk_layout.end_idx_exclusive = int(blk_layout.global_end_idx) + 1
+            if blk_layout.start_idx is None:
+                blk_layout.start_idx = int(blk_layout.global_start_idx)
+            assert_layout_specials_complete(
+                block_layout=blk_layout,
+                lead_names=ecg_blocks[blk_idx]["lead_names"],
+            )
+            all_specials = []
+            if blk_layout.global_start_idx is not None:
+                all_specials.append(blk_layout.global_start_idx)
+            all_specials.extend(list(blk_layout.lead_start_idx.values()))
+            all_specials.extend(list(blk_layout.lead_end_idx.values()))
+            if blk_layout.global_end_idx is not None:
+                all_specials.append(blk_layout.global_end_idx)
+            blk_layout.special_idxs_sorted = sorted(all_specials)
+            blk_layout.signal_pos_list = sorted(
+                [p for lst in blk_layout.signal_pos_by_lead.values() for p in lst]
+            )
+
+        mask_result = self.mask_strategy.build(
+            layout,
+            device=device,
+            dtype=inputs_embeds.dtype,
+        )
+        use_return = self.packing_schema.conversation.format_id == "harmony_chat_v1"
+        _, stop_text = turn_wrappers(self.packing_schema, prompt_tokens.model_role, use_return=use_return)
+        stop_ids = tok.encode(stop_text, add_special_tokens=False)
+        return PromptContext(
+            inputs_embeds=inputs_embeds,
+            layout=layout,
+            prompt_preview=prompt_preview,
+            stop_ids=stop_ids,
+            input_embedder=input_embedder,
+            mask_strategy=self.mask_strategy,
+            mask_result=mask_result,
+        )
+
+    def _autoregressive_generate(
+        self,
+        *,
+        context: PromptContext,
+        max_new_tokens: int,
+        temperature: float,
+        top_k: Optional[int],
+        top_p: float,
+        min_p: float,
+    ) -> List[int]:
+        tok = self.tokenizer
+        wrapper = self.wrapper
+        device = context.inputs_embeds.device
+        embeds = context.inputs_embeds.clone()
+        layout = context.layout
+        input_embedder = context.input_embedder
+        mask_result = context.mask_result
+
+        generated: List[int] = []
+        stop_ids = context.stop_ids
+        stop_len = len(stop_ids)
+        eos_id = tok.eos_token_id
+
+        for _ in range(max_new_tokens):
+            additive = mask_result.additive.unsqueeze(0).unsqueeze(0)
+            outputs = wrapper.forward_language_model(
+                inputs_embeds=embeds.unsqueeze(0),
+                attention_mask=additive,
+                labels=None,
+            )
+            logits = outputs.logits[0, -1, :].float()
+            next_token = _sample_next_token(
+                logits,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                min_p=min_p,
+            )
+            token_id = int(next_token.item())
+            generated.append(token_id)
+
+            if stop_len and generated[-stop_len:] == stop_ids:
+                generated = generated[:-stop_len]
+                break
+            if eos_id is not None and token_id == eos_id:
+                break
+
+            with torch.no_grad():
+                new_embed = wrapper.tokens_to_embeds(input_embedder, [token_id], device=device)
+            embeds = torch.cat([embeds, new_embed], dim=0)
+
+            layout.seq_len = embeds.size(0)
+            new_idx = layout.seq_len - 1
+            layout.text_idxs.append(new_idx)
+            mask_result = context.mask_strategy.update_for_generated_token(
+                layout,
+                device=device,
+                dtype=embeds.dtype,
+                previous=mask_result,
+            )
+            context.mask_result = mask_result
+
+        return generated

camel_inference/src/camel/model_init.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Tuple, Type
+import torch
+import torch.nn as nn
+from peft import (
+    LoraConfig,
+    get_peft_model,
+    TaskType,
+)
+
+from camel.ecg_gemma_model import ECGGemmaPrefix as ECGModelPrefix
+
+def attach_lora(
+    model: nn.Module,
+    lora_cfg_dict: Dict[str, Any],
+    device: torch.device,
+) -> Tuple[nn.Module, LoraConfig]:
+    """Attach LoRA adapters to the frozen model, leaving only LoRA trainable."""
+    cfg = LoraConfig(
+        r=int(lora_cfg_dict["r"]),
+        lora_alpha=int(lora_cfg_dict.get("lora_alpha", int(lora_cfg_dict["r"]) * 2)),
+        lora_dropout=float(lora_cfg_dict.get("lora_dropout", 0.0)),
+        target_modules=list(lora_cfg_dict.get("target_modules", [])),
+        task_type=TaskType(lora_cfg_dict.get("task_type", "CAUSAL_LM")),
+        bias=lora_cfg_dict.get("bias", "none"),
+        inference_mode=False,
+        use_dora=bool(lora_cfg_dict.get("use_dora", False)),
+    )
+    model = get_peft_model(model, cfg)
+    model.to(device)
+    return model, cfg
+
+def build_conv_encoder(
+    *,
+    conv_ckpt_path: Optional[str],
+    device: torch.device,
+    unfreeze: bool = False,
+) -> nn.Module:
+    """
+    Build the 1D conv stack and load weights from the provided checkpoint, including
+    key normalization and optional unfreezing. Identical to train_ecg_text.py.
+    """
+    enc = nn.Sequential(
+        nn.Conv1d(1,   32, kernel_size=4, stride=2, padding=1),  # L:256->128
+        nn.ReLU(inplace=True),
+        nn.Conv1d(32,  64, kernel_size=4, stride=2, padding=1),  # L:128->64
+        nn.ReLU(inplace=True),
+        nn.Conv1d(64, 128, kernel_size=4, stride=2, padding=1),  # L:64->32
+        nn.ReLU(inplace=True),
+        nn.Conv1d(128, 4,  kernel_size=4, stride=2, padding=1),  # L:32->16, C:4
+        nn.ReLU(inplace=True),
+    ).to(device=device, dtype=torch.float32)
+
+    if conv_ckpt_path:
+        ckpt = torch.load(conv_ckpt_path, map_location="cpu")
+        raw_sd = ckpt["model_state_dict"] if isinstance(ckpt, dict) and "model_state_dict" in ckpt else ckpt
+        norm_sd: Dict[str, torch.Tensor] = {}
+        for k, v in raw_sd.items():
+            kk = k
+            if kk.startswith("module."):
+                kk = kk[len("module."):]
+            if kk.startswith("_orig_mod."):
+                kk = kk[len("_orig_mod."):]
+            if kk.startswith("enc."):
+                kk = kk[len("enc."):]
+            norm_sd[kk] = v
+        wanted = {f"{i}.{w}" for i in (0, 2, 4, 6) for w in ("weight", "bias")}
+        conv_sd = {k: v for k, v in norm_sd.items() if k in wanted}
+        missing, unexpected = enc.load_state_dict(conv_sd, strict=True)
+        if missing or unexpected:
+            print(f"[conv load] missing={list(missing)} unexpected={list(unexpected)}")
+    enc.eval()
+    return enc
+
+def build_wrapper(
+    *,
+    wrapper_cls: Type[nn.Module] = ECGModelPrefix,
+    language_model: nn.Module,
+    conv_encoder: nn.Module,
+    hidden_size: int,
+    num_ecg_special_tokens: int,
+    dtype: torch.dtype,
+    enc_out_dim: int = 64,
+    freeze_encoder: bool = True,
+    inference: bool = False,
+    projector_name: str = "linear",
+) -> ECGModelPrefix:
+    """Construct the ECG-language wrapper (keeps default wrapper class)."""
+    wrapper = wrapper_cls(
+        language_model,
+        enc=conv_encoder,
+        hidden_size=hidden_size,
+        num_ecg_special_tokens=num_ecg_special_tokens,
+        dtype=dtype,
+        enc_out_dim=enc_out_dim,
+        freeze_encoder=freeze_encoder,
+        inference=inference,
+        projector_name=projector_name,
+    )
+    return wrapper
+
+
+__all__ = [
+    "attach_lora",
+    "build_conv_encoder",
+    "build_wrapper",
+]

camel_inference/src/camel/model_introspect.py

@@ -0,0 +1,95 @@
+"""Model introspection utilities driven by registry hints.
+
+Centralizes resolution of model-internal structures to avoid hardcoded
+attribute names in call sites. Use the hint paths defined in the model
+registry to locate transformer layers and config attributes.
+"""
+from __future__ import annotations
+
+from typing import List, Optional, Sequence
+import torch.nn as nn
+
+def _walk_attr_path(root: object, dotted_path: str) -> Optional[object]:
+    cur: object = root
+    for part in dotted_path.split("."):
+        if not hasattr(cur, part):
+            return None
+        cur = getattr(cur, part)
+    return cur
+
+def resolve_layers(model: nn.Module, path_hints: Sequence[str]) -> List[nn.Module]:
+    """
+    Resolve the text transformer layer sequence using the first successful
+    dotted path from `path_hints` relative to common roots.
+
+    We try against several candidate roots to be robust to wrappers (e.g., PEFT):
+      - the model itself
+      - model.base_model (if present)
+      - model.base_model.model (if present)
+
+    We also try a small set of generic fallback hints ("model.language_model.layers",
+    "language_model.layers", "model.layers", "layers") if the provided hints fail.
+    """
+    roots: List[object] = [model]
+    base = getattr(model, "base_model", None)
+    if base is not None:
+        roots.append(base)
+        base_model_attr = getattr(base, "model", None)
+        if base_model_attr is not None:
+            roots.append(base_model_attr)
+
+    tried: List[str] = []
+    def _try_hints(root: object, hints: Sequence[str]) -> Optional[List[nn.Module]]:
+        for hint in hints:
+            tried.append(hint)
+            obj = _walk_attr_path(root, hint)
+            if obj is None:
+                continue
+            if isinstance(obj, (list, tuple)) and all(isinstance(x, nn.Module) for x in obj):
+                return list(obj)
+            if hasattr(obj, "__iter__") and not isinstance(obj, (str, bytes)):
+                try:
+                    seq = list(obj)
+                    if seq and all(isinstance(x, nn.Module) for x in seq):
+                        return seq
+                except Exception:
+                    pass
+        return None
+
+    # Try provided hints against all roots
+    for root in roots:
+        found = _try_hints(root, path_hints)
+        if found is not None:
+            return found
+
+    # Fallback generic hints against all roots
+    generic_hints = (
+        "model.language_model.layers",
+        "language_model.layers",
+        "model.layers",
+        "layers",
+    )
+    for root in roots:
+        found = _try_hints(root, generic_hints)
+        if found is not None:
+            return found
+
+    raise RuntimeError(
+        f"Could not resolve transformer layers via provided hints: {list(path_hints)}"
+    )
+
+def resolve_hidden_size(model: nn.Module, attr_paths: Sequence[str]) -> int:
+    """Resolve hidden size via the first successful dotted config attribute path."""
+    for path in attr_paths:
+        val = _walk_attr_path(model, path)
+        if isinstance(val, (int, float)):
+            return int(val)
+    raise AttributeError(
+        f"Could not resolve hidden size from any of: {list(attr_paths)}"
+    )
+
+
+__all__ = [
+    "resolve_layers",
+    "resolve_hidden_size",
+]

camel_inference/src/camel/model_registry.py

@@ -0,0 +1,497 @@
+"""
+Utilities for loading per-model configuration metadata used across training and inference.
+
+The registry is defined in YAML (see model_registry.yaml in this directory) and exposes
+immutable ModelConfig objects for downstream consumers. The intent is to centralize
+model-specific defaults (prompt format, tokenizer quirks, wrapper class path, LoRA constraints, etc.)
+so that adding support for a new backbone primarily involves updating the registry.
+"""
+from __future__ import annotations
+
+import copy
+import importlib
+import dataclasses
+import os
+from collections.abc import Mapping as ABCMapping, Sequence as ABCSequence
+from pathlib import Path
+from types import MappingProxyType
+from typing import Any, Dict, Iterable, Mapping, Optional, Tuple
+import yaml
+
+class ModelRegistryError(RuntimeError):
+    """Raised when the registry file is missing or malformed."""
+
+@dataclasses.dataclass(frozen=True)
+class PromptConfig:
+    start_of_turn: str
+    end_of_turn: str
+    roles: Mapping[str, str]
+    enforce_bos: bool
+    enforce_eos: bool
+    allow_multiple_eos: bool
+
+
+@dataclasses.dataclass(frozen=True)
+class TokenizerConfig:
+    pad_token_strategy: str
+    require_bos: bool
+    require_eos: bool
+    use_fast: bool = True
+    add_prefix_space: bool = False
+
+
+@dataclasses.dataclass(frozen=True)
+class ArchitectureConfig:
+    wrapper_class: str
+    hidden_size_attrs: Tuple[str, ...]
+    language_model_path_hints: Tuple[str, ...]
+    attn_implementation: str
+    conv_out_dim: Optional[int] = None
+
+
+@dataclasses.dataclass(frozen=True)
+class LoRAPolicyConfig:
+    expect_language_only: bool
+    allowed_markers: Tuple[str, ...]
+    blocked_markers: Tuple[str, ...]
+    freeze_vision: bool
+
+
+@dataclasses.dataclass(frozen=True)
+class PackingConversationConfig:
+    format_id: str
+    user_role_aliases: Tuple[str, ...]
+    model_role_aliases: Tuple[str, ...]
+    strip_image_from_roles: Tuple[str, ...]
+    merge_system_with_first_user: bool
+
+
+@dataclasses.dataclass(frozen=True)
+class PackingECGTokensConfig:
+    global_start: str
+    global_end: str
+    lead_start_template: str
+    lead_end_template: str
+    canonical_leads: Tuple[str, ...]
+
+
+@dataclasses.dataclass(frozen=True)
+class PackingConfig:
+    prompt_format: str
+    conversation: PackingConversationConfig
+    ecg_tokens: PackingECGTokensConfig
+
+
+@dataclasses.dataclass(frozen=True)
+class ModelConfig:
+    """Typed wrapper over a single model entry in the registry."""
+
+    name: str
+    data: Mapping[str, Any]
+
+    @property
+    def hf_id(self) -> str:
+        return _require_str(self.data, "hf_id", self.name)
+
+    @property
+    def prompt(self) -> Mapping[str, Any]:
+        return _require_mapping(self.data, "prompt", self.name)
+
+    @property
+    def tokenizer(self) -> Mapping[str, Any]:
+        return _require_mapping(self.data, "tokenizer", self.name)
+
+    @property
+    def architecture(self) -> Mapping[str, Any]:
+        return _require_mapping(self.data, "architecture", self.name)
+
+    @property
+    def lora_policy(self) -> Mapping[str, Any]:
+        return _require_mapping(self.data, "lora_policy", self.name)
+
+    @property
+    def packing(self) -> Mapping[str, Any]:
+        return _require_mapping(self.data, "packing", self.name)
+
+    def prompt_config(self) -> PromptConfig:
+        prompt = self.prompt
+        roles = _require_mapping(prompt, "roles", self.name, section="prompt")
+        return PromptConfig(
+            start_of_turn=_require_str(prompt, "start_of_turn", self.name, section="prompt"),
+            end_of_turn=_require_str(prompt, "end_of_turn", self.name, section="prompt"),
+            roles={k: str(v) for k, v in roles.items()},
+            enforce_bos=_require_bool(prompt, "enforce_bos", self.name, section="prompt"),
+            enforce_eos=_require_bool(prompt, "enforce_eos", self.name, section="prompt"),
+            allow_multiple_eos=_optional_bool(prompt, "allow_multiple_eos", False, self.name, section="prompt"),
+        )
+
+    def required_prompt_text(self, key: str) -> str:
+        return _require_str(self.prompt, key, self.name, section="prompt")
+
+    def tokenizer_config(self) -> TokenizerConfig:
+        tokenizer = self.tokenizer
+        pad_strategy = _require_str(tokenizer, "pad_token_strategy", self.name, section="tokenizer")
+        require_bos = _require_bool(tokenizer, "require_bos", self.name, section="tokenizer")
+        require_eos = _require_bool(tokenizer, "require_eos", self.name, section="tokenizer")
+        # Optional fields with defaults for backward compatibility
+        use_fast = _optional_bool(tokenizer, "use_fast", True, self.name, section="tokenizer")
+        add_prefix_space = _optional_bool(tokenizer, "add_prefix_space", False, self.name, section="tokenizer")
+        return TokenizerConfig(
+            pad_token_strategy=pad_strategy,
+            require_bos=require_bos,
+            require_eos=require_eos,
+            use_fast=use_fast,
+            add_prefix_space=add_prefix_space,
+        )
+
+    def architecture_config(self) -> ArchitectureConfig:
+        arch = self.architecture
+        conv_out = arch.get("conv_out_dim") if isinstance(arch, ABCMapping) else None
+        try:
+            conv_out_int = int(conv_out) if conv_out is not None else None
+        except Exception:
+            conv_out_int = None
+        return ArchitectureConfig(
+            wrapper_class=_require_str(arch, "wrapper_class", self.name, section="architecture"),
+            hidden_size_attrs=tuple(
+                _require_sequence_of_str(arch, "hidden_size_attrs", self.name, section="architecture")
+            ),
+            language_model_path_hints=tuple(
+                _require_sequence_of_str(arch, "language_model_path_hints", self.name, section="architecture")
+            ),
+            attn_implementation=_require_str(arch, "attn_implementation", self.name, section="architecture"),
+            conv_out_dim=conv_out_int,
+        )
+
+    def lora_policy_config(self) -> LoRAPolicyConfig:
+        lora = self.lora_policy
+        return LoRAPolicyConfig(
+            expect_language_only=_require_bool(lora, "expect_language_only", self.name, section="lora_policy"),
+            allowed_markers=tuple(
+                _require_sequence_of_str(lora, "allowed_markers", self.name, section="lora_policy")
+            ),
+            blocked_markers=tuple(
+                _require_sequence_of_str(lora, "blocked_markers", self.name, section="lora_policy")
+            ),
+            freeze_vision=_require_bool(lora, "freeze_vision", self.name, section="lora_policy"),
+        )
+
+    def packing_config(self) -> PackingConfig:
+        packing = self.packing
+        format_id = _require_str(packing, "prompt_format", self.name, section="packing")
+
+        conversation = _require_mapping(packing, "conversation", self.name, section="packing")
+        user_aliases = tuple(
+            _require_sequence_of_str(conversation, "user_role_aliases", self.name, section="packing.conversation")
+        )
+        model_aliases = tuple(
+            _require_sequence_of_str(conversation, "model_role_aliases", self.name, section="packing.conversation")
+        )
+        strip_roles = tuple(
+            _require_sequence_of_str(conversation, "strip_image_from_roles", self.name, section="packing.conversation")
+        )
+        merge_system = _require_bool(
+            conversation, "merge_system_with_first_user", self.name, section="packing.conversation"
+        )
+        conv_cfg = PackingConversationConfig(
+            format_id=format_id,
+            user_role_aliases=user_aliases,
+            model_role_aliases=model_aliases,
+            strip_image_from_roles=strip_roles,
+            merge_system_with_first_user=merge_system,
+        )
+
+        ecg_tokens = _require_mapping(packing, "ecg_tokens", self.name, section="packing")
+        global_start = _require_str(ecg_tokens, "global_start", self.name, section="packing.ecg_tokens")
+        global_end = _require_str(ecg_tokens, "global_end", self.name, section="packing.ecg_tokens")
+        lead_start_template = _require_str(ecg_tokens, "lead_start_template", self.name, section="packing.ecg_tokens")
+        lead_end_template = _require_str(ecg_tokens, "lead_end_template", self.name, section="packing.ecg_tokens")
+        canonical_leads = tuple(
+            _require_sequence_of_str(ecg_tokens, "canonical_leads", self.name, section="packing.ecg_tokens")
+        )
+        ecg_cfg = PackingECGTokensConfig(
+            global_start=global_start,
+            global_end=global_end,
+            lead_start_template=lead_start_template,
+            lead_end_template=lead_end_template,
+            canonical_leads=canonical_leads,
+        )
+
+        return PackingConfig(
+            prompt_format=format_id,
+            conversation=conv_cfg,
+            ecg_tokens=ecg_cfg,
+        )
+
+    def resolve_wrapper_class(self):
+        arch = self.architecture_config()
+        if "." not in arch.wrapper_class:
+            raise ModelRegistryError(
+                f"Wrapper class path '{arch.wrapper_class}' for model '{self.name}' must be in 'module.ClassName' form."
+            )
+        module_name, class_name = arch.wrapper_class.rsplit(".", 1)
+        try:
+            module = importlib.import_module("." + module_name, package=__package__)
+        except ImportError as exc:
+            raise ModelRegistryError(
+                f"Failed to import wrapper module '{module_name}' for model '{self.name}': {exc}"
+            ) from exc
+        try:
+            wrapper_cls = getattr(module, class_name)
+        except AttributeError as exc:
+            raise ModelRegistryError(
+                f"Wrapper class '{class_name}' not found in module '{module_name}' for model '{self.name}'."
+            ) from exc
+        return wrapper_cls
+
+
+def _default_registry_path() -> Path:
+    return Path(__file__).resolve().with_name("model_registry.yaml")
+
+
+def load_registry(
+    *,
+    registry_path: Optional[os.PathLike[str] | str] = None,
+    model_overrides: Optional[Mapping[str, Mapping[str, Any]]] = None,
+) -> "ModelRegistry":
+    """
+    Load the model registry from YAML.
+
+    Args:
+        registry_path: Optional path to the YAML file. Defaults to `model_registry.yaml` alongside this module.
+        model_overrides: Optional mapping of model name -> override dict that will be deep-merged
+                         onto the YAML entry (useful for ad-hoc experimentation).
+    """
+    path = Path(registry_path) if registry_path is not None else _default_registry_path()
+    if not path.exists():
+        raise ModelRegistryError(f"Model registry file not found at {path}")
+    try:
+        with path.open("r", encoding="utf-8") as fh:
+            raw = yaml.safe_load(fh)
+    except yaml.YAMLError as exc:
+        raise ModelRegistryError(f"Failed to parse model registry YAML: {exc}") from exc
+
+    if not isinstance(raw, ABCMapping):
+        raise ModelRegistryError("Model registry root must be a mapping.")
+
+    models = raw.get("models")
+    if not isinstance(models, ABCMapping) or len(models) == 0:
+        raise ModelRegistryError("Model registry must define a non-empty 'models' mapping.")
+
+    entries: Dict[str, Mapping[str, Any]] = {}
+    for name, cfg in models.items():
+        if not isinstance(name, str):
+            raise ModelRegistryError("Model names must be strings.")
+        if not isinstance(cfg, ABCMapping):
+            raise ModelRegistryError(f"Model '{name}' entry must be a mapping.")
+        merged = copy.deepcopy(cfg)
+        if model_overrides and name in model_overrides:
+            _deep_update(merged, model_overrides[name])
+        _validate_model_entry(name, merged)
+        entries[name] = merged
+
+    return ModelRegistry(entries, source_path=path)
+
+
+class ModelRegistry:
+    """In-memory view over the registry."""
+
+    def __init__(self, models: Mapping[str, Mapping[str, Any]], *, source_path: Path):
+        self._models = dict(models)
+        self._source_path = Path(source_path)
+
+    @property
+    def source_path(self) -> Path:
+        return self._source_path
+
+    def names(self) -> Iterable[str]:
+        return tuple(self._models.keys())
+
+    def get(self, name: str) -> ModelConfig:
+        if name not in self._models:
+            raise ModelRegistryError(f"Unknown model '{name}'. Known models: {sorted(self._models)}")
+        return ModelConfig(name=name, data=_deep_freeze(self._models[name]))
+
+
+def _validate_model_entry(name: str, entry: Mapping[str, Any]) -> None:
+    # Required string
+    _require_str(entry, "hf_id", name)
+
+    # Prompt section
+    prompt = _require_mapping(entry, "prompt", name)
+    _require_str(prompt, "start_of_turn", name, section="prompt")
+    _require_str(prompt, "end_of_turn", name, section="prompt")
+    roles = _require_mapping(prompt, "roles", name, section="prompt")
+    for role_key in ("user", "model"):
+        _require_str(roles, role_key, name, section="prompt.roles")
+    for flag in ("enforce_bos", "enforce_eos"):
+        _require_bool(prompt, flag, name, section="prompt")
+
+    # Tokenizer section
+    tokenizer = _require_mapping(entry, "tokenizer", name)
+    _require_str(tokenizer, "pad_token_strategy", name, section="tokenizer")
+    for flag in ("require_bos", "require_eos"):
+        _require_bool(tokenizer, flag, name, section="tokenizer")
+
+    # Architecture
+    architecture = _require_mapping(entry, "architecture", name)
+    _require_str(architecture, "wrapper_class", name, section="architecture")
+    hidden_attrs = _require_sequence_of_str(architecture, "hidden_size_attrs", name, section="architecture")
+    if len(hidden_attrs) == 0:
+        raise ModelRegistryError(
+            f"Model '{name}' architecture.hidden_size_attrs must contain at least one attribute path."
+        )
+    _require_sequence_of_str(architecture, "language_model_path_hints", name, section="architecture")
+    _require_str(architecture, "attn_implementation", name, section="architecture")
+    # conv_out_dim is optional; if present, ensure it is an int-like
+    if "conv_out_dim" in architecture:
+        val = architecture.get("conv_out_dim")
+        try:
+            int(val)  # type: ignore[arg-type]
+        except Exception:
+            raise ModelRegistryError(
+                f"Model '{name}' field 'architecture.conv_out_dim' must be an integer when provided."
+            )
+
+    # LoRA policy
+    lora_policy = _require_mapping(entry, "lora_policy", name)
+    _require_bool(lora_policy, "expect_language_only", name, section="lora_policy")
+    allowed = _require_sequence_of_str(lora_policy, "allowed_markers", name, section="lora_policy")
+    blocked = _require_sequence_of_str(lora_policy, "blocked_markers", name, section="lora_policy")
+    overlap = set(allowed).intersection(blocked)
+    if overlap:
+        raise ModelRegistryError(
+            f"Model '{name}' lora_policy.allowed_markers and lora_policy.blocked_markers overlap: {sorted(overlap)}"
+        )
+    _require_bool(lora_policy, "freeze_vision", name, section="lora_policy")
+
+    # Packing
+    packing = _require_mapping(entry, "packing", name)
+    _require_str(packing, "prompt_format", name, section="packing")
+    conversation = _require_mapping(packing, "conversation", name, section="packing")
+    _require_sequence_of_str(conversation, "user_role_aliases", name, section="packing.conversation")
+    _require_sequence_of_str(conversation, "model_role_aliases", name, section="packing.conversation")
+    _require_sequence_of_str(conversation, "strip_image_from_roles", name, section="packing.conversation")
+    _require_bool(conversation, "merge_system_with_first_user", name, section="packing.conversation")
+    ecg_tokens = _require_mapping(packing, "ecg_tokens", name, section="packing")
+    _require_str(ecg_tokens, "global_start", name, section="packing.ecg_tokens")
+    _require_str(ecg_tokens, "global_end", name, section="packing.ecg_tokens")
+    _require_str(ecg_tokens, "lead_start_template", name, section="packing.ecg_tokens")
+    _require_str(ecg_tokens, "lead_end_template", name, section="packing.ecg_tokens")
+    if len(_require_sequence_of_str(ecg_tokens, "canonical_leads", name, section="packing.ecg_tokens")) == 0:
+        raise ModelRegistryError(
+            f"Model '{name}' packing.ecg_tokens.canonical_leads must contain at least one lead."
+        )
+
+
+def _require_mapping(
+    parent: Mapping[str, Any],
+    key: str,
+    model_name: str,
+    *,
+    section: Optional[str] = None,
+) -> Mapping[str, Any]:
+    value = parent.get(key)
+    if not isinstance(value, ABCMapping):
+        loc = f"{section}.{key}" if section else key
+        raise ModelRegistryError(f"Model '{model_name}' is missing mapping '{loc}'.")
+    return value
+
+
+def _require_str(
+    parent: Mapping[str, Any],
+    key: str,
+    model_name: str,
+    *,
+    section: Optional[str] = None,
+) -> str:
+    value = parent.get(key)
+    if not isinstance(value, str):
+        loc = f"{section}.{key}" if section else key
+        raise ModelRegistryError(f"Model '{model_name}' field '{loc}' must be a string.")
+    return value
+
+
+def _require_bool(
+    parent: Mapping[str, Any],
+    key: str,
+    model_name: str,
+    *,
+    section: Optional[str] = None,
+) -> bool:
+    value = parent.get(key)
+    if not isinstance(value, bool):
+        loc = f"{section}.{key}" if section else key
+        raise ModelRegistryError(f"Model '{model_name}' field '{loc}' must be a boolean.")
+    return value
+
+
+def _require_sequence_of_str(
+    parent: Mapping[str, Any],
+    key: str,
+    model_name: str,
+    *,
+    section: Optional[str] = None,
+) -> Tuple[str, ...]:
+    value = parent.get(key)
+    if not isinstance(value, ABCSequence) or isinstance(value, (str, bytes)):
+        loc = f"{section}.{key}" if section else key
+        raise ModelRegistryError(f"Model '{model_name}' field '{loc}' must be a sequence of strings.")
+    items = []
+    for item in value:
+        if not isinstance(item, str):
+            loc = f"{section}.{key}" if section else key
+            raise ModelRegistryError(f"Model '{model_name}' field '{loc}' must contain only strings.")
+        items.append(item)
+    return tuple(items)
+
+
+def _deep_update(target: Dict[str, Any], updates: Mapping[str, Any]) -> None:
+    """
+    Recursively merge `updates` into `target` in-place.
+    """
+    for key, value in updates.items():
+        if isinstance(value, ABCMapping) and isinstance(target.get(key), dict):
+            _deep_update(target[key], value)  # type: ignore[arg-type]
+        else:
+            target[key] = copy.deepcopy(value)
+
+
+def _optional_bool(
+    parent: Mapping[str, Any],
+    key: str,
+    default: bool,
+    model_name: str,
+    *,
+    section: Optional[str] = None,
+) -> bool:
+    if key not in parent:
+        return default
+    value = parent.get(key)
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, int) and value in (0, 1):
+        return bool(value)
+    if isinstance(value, str):
+        normalized = value.strip().lower()
+        if normalized in {"true", "yes", "y", "1"}:
+            return True
+        if normalized in {"false", "no", "n", "0"}:
+            return False
+    loc = f"{section}.{key}" if section else key
+    raise ModelRegistryError(f"Model '{model_name}' field '{loc}' must be a boolean when provided.")
+
+
+def _deep_freeze(obj: Any) -> Any:
+    """
+    Recursively convert mutable containers to immutable/read-only equivalents.
+    """
+    if isinstance(obj, dict):
+        return MappingProxyType({k: _deep_freeze(v) for k, v in obj.items()})
+    if isinstance(obj, list):
+        return tuple(_deep_freeze(v) for v in obj)
+    if isinstance(obj, tuple):
+        return tuple(_deep_freeze(v) for v in obj)
+    if isinstance(obj, set):
+        return frozenset(_deep_freeze(v) for v in obj)
+    return obj

camel_inference/src/camel/model_registry.yaml

@@ -0,0 +1,509 @@
+# Model Registry
+#
+# This file defines per-backbone configuration consumed by training and inference.
+# Fields overview:
+# - hf_id: Hugging Face model identifier for the frozen language model.
+# - prompt: Chat formatting tokens and roles used to build prompts
+#   - start_of_turn / end_of_turn: literal strings delimiting speaker turns
+#   - roles.user / roles.model: canonical role names for user/model
+#   - enforce_bos / enforce_eos: whether BOS/EOS must appear at the start/end
+# - tokenizer:
+#   - pad_token_strategy: how to set pad token ("eos" uses eos as pad)
+#   - require_bos / require_eos: tokenizer must expose these tokens
+# - architecture:
+#   - wrapper_class: import path to the ECG wrapper class (module.Class)
+#   - hidden_size_attrs: ordered config attribute paths to read hidden size
+#   - language_model_path_hints: ordered attribute paths to locate transformer layers
+#   - attn_implementation: preferred attention backend to use when loading
+# - lora_policy:
+#   - expect_language_only: LoRA must live under language/text model stacks
+#   - allowed_markers / blocked_markers: substrings used to validate LoRA placement
+#   - freeze_vision: freeze LoRA params under blocked stacks during training
+# - packing:
+#   - prompt_format: prompt/conversation template id
+#   - conversation: conversation role normalization and preprocessing
+#   - ecg_tokens: special tokens inserted to mark ECG structure
+models:
+  medgemma-27b-it:
+    hf_id: "google/medgemma-27b-text-it"
+    prompt:
+      start_of_turn: "<start_of_turn>"
+      end_of_turn: "<end_of_turn>\n"
+      roles:
+        user: "user"
+        model: "model"
+      enforce_bos: true
+      enforce_eos: true
+    tokenizer:
+      pad_token_strategy: "eos"
+      require_bos: true
+      require_eos: true
+      use_fast: true
+      add_prefix_space: false
+    architecture:
+      wrapper_class: "ecg_gemma_model.ECGGemmaPrefix"
+      hidden_size_attrs:
+        - "config.hidden_size"
+        - "config.text_config.hidden_size"
+      language_model_path_hints:
+        - "base_model.model.language_model.layers"
+        - "model.language_model.layers"
+        - "model.layers"
+      attn_implementation: "eager"
+      conv_out_dim: 64
+    lora_policy:
+      expect_language_only: true
+      allowed_markers:
+        - "language_model"
+        - ".model.layers."
+        - ".layers."
+      blocked_markers:
+        - "vision"
+        - "multi_modal"
+        - "projector"
+        - ".enc."
+        - "encoder_proj"
+      freeze_vision: true
+    packing:
+      prompt_format: "gemma_chat_v1"
+      conversation:
+        user_role_aliases: ["human", "user"]
+        model_role_aliases: ["gpt", "assistant"]
+        strip_image_from_roles: ["human"]
+        merge_system_with_first_user: true
+      ecg_tokens:
+        global_start: "<ecg_global_start>"
+        global_end: "<ecg_global_end>"
+        lead_start_template: "<ecg_lead_{lead_lower}_start>"
+        lead_end_template: "<ecg_lead_{lead_lower}_end>"
+        canonical_leads: ["I", "II", "III", "aVR", "aVL", "aVF", "V1", "V2", "V3", "V4", "V5", "V6"]
+  gemma-12b-it:
+    hf_id: "google/gemma-3-12b-it"
+    prompt:
+      start_of_turn: "<start_of_turn>"
+      end_of_turn: "<end_of_turn>\n"
+      roles:
+        user: "user"
+        model: "model"
+      enforce_bos: true
+      enforce_eos: true
+    tokenizer:
+      pad_token_strategy: "eos"
+      require_bos: true
+      require_eos: true
+      use_fast: true
+      add_prefix_space: false
+    architecture:
+      wrapper_class: "ecg_gemma_model.ECGGemmaPrefix"
+      hidden_size_attrs:
+        - "config.hidden_size"
+        - "config.text_config.hidden_size"
+      language_model_path_hints:
+        - "base_model.model.language_model.layers"
+        - "model.language_model.layers"
+        - "model.layers"
+      attn_implementation: "eager"
+      conv_out_dim: 64
+    lora_policy:
+      expect_language_only: true
+      allowed_markers:
+        - "language_model"
+        - ".model.layers."
+        - ".layers."
+      blocked_markers:
+        - "vision"
+        - "multi_modal"
+        - "projector"
+        - ".enc."
+        - "encoder_proj"
+      freeze_vision: true
+    packing:
+      prompt_format: "gemma_chat_v1"
+      conversation:
+        user_role_aliases: ["human", "user"]
+        model_role_aliases: ["gpt", "assistant"]
+        strip_image_from_roles: ["human"]
+        merge_system_with_first_user: true
+      ecg_tokens:
+        global_start: "<ecg_global_start>"
+        global_end: "<ecg_global_end>"
+        lead_start_template: "<ecg_lead_{lead_lower}_start>"
+        lead_end_template: "<ecg_lead_{lead_lower}_end>"
+        canonical_leads: ["I", "II", "III", "aVR", "aVL", "aVF", "V1", "V2", "V3", "V4", "V5", "V6"]
+  medgemma-4b-it:
+    hf_id: "google/medgemma-4b-it"
+    prompt:
+      start_of_turn: "<start_of_turn>"
+      end_of_turn: "<end_of_turn>\n"
+      roles:
+        user: "user"
+        model: "model"
+      enforce_bos: true
+      enforce_eos: true
+    tokenizer:
+      pad_token_strategy: "eos"
+      require_bos: true
+      require_eos: true
+      use_fast: true
+      add_prefix_space: false
+    architecture:
+      wrapper_class: "ecg_gemma_model.ECGGemmaPrefix"
+      hidden_size_attrs:
+        - "config.hidden_size"
+        - "config.text_config.hidden_size"
+      language_model_path_hints:
+        - "base_model.model.language_model.layers"
+        - "model.language_model.layers"
+        - "model.layers"
+      attn_implementation: "eager"
+      conv_out_dim: 64
+    lora_policy:
+      expect_language_only: true
+      allowed_markers:
+        - "language_model"
+        - ".model.layers."
+        - ".layers."
+      blocked_markers:
+        - "vision"
+        - "multi_modal"
+        - "projector"
+        - ".enc."
+        - "encoder_proj"
+      freeze_vision: true
+    packing:
+      prompt_format: "gemma_chat_v1"
+      conversation:
+        user_role_aliases: ["human", "user"]
+        model_role_aliases: ["gpt", "assistant"]
+        strip_image_from_roles: ["human"]
+        merge_system_with_first_user: true
+      ecg_tokens:
+        global_start: "<ecg_global_start>"
+        global_end: "<ecg_global_end>"
+        lead_start_template: "<ecg_lead_{lead_lower}_start>"
+        lead_end_template: "<ecg_lead_{lead_lower}_end>"
+        canonical_leads: ["I", "II", "III", "aVR", "aVL", "aVF", "V1", "V2", "V3", "V4", "V5", "V6"]
+  qwen3-4b-instruct:
+    hf_id: "Qwen/Qwen3-4B-Instruct-2507"
+    prompt:
+      start_of_turn: "<|im_start|>"
+      end_of_turn: "<|im_end|>"
+      roles:
+        user: "user"
+        model: "assistant"
+      enforce_bos: false
+      enforce_eos: false
+      allow_multiple_eos: true
+    tokenizer:
+      pad_token_strategy: "existing"
+      require_bos: false
+      require_eos: false
+      use_fast: true
+      add_prefix_space: false
+    architecture:
+      wrapper_class: "ecg_qwen_model.ECGQwenPrefix"
+      hidden_size_attrs:
+        - "config.hidden_size"
+      language_model_path_hints:
+        - "model.layers"
+        - "model.model.layers"
+      attn_implementation: "eager"
+      conv_out_dim: 64
+    lora_policy:
+      expect_language_only: true
+      allowed_markers:
+        - "model.layers."
+        - ".layers."
+      blocked_markers:
+        - "vision"
+        - "multi_modal"
+        - "projector"
+        - ".enc."
+        - "encoder_proj"
+      freeze_vision: false
+    packing:
+      prompt_format: "qwen_chat_v1"
+      conversation:
+        user_role_aliases: ["human", "user"]
+        model_role_aliases: ["gpt", "assistant"]
+        strip_image_from_roles: ["human"]
+        merge_system_with_first_user: true
+      ecg_tokens:
+        global_start: "<ecg_global_start>"
+        global_end: "<ecg_global_end>"
+        lead_start_template: "<ecg_lead_{lead_lower}_start>"
+        lead_end_template: "<ecg_lead_{lead_lower}_end>"
+        canonical_leads: ["I", "II", "III", "aVR", "aVL", "aVF", "V1", "V2", "V3", "V4", "V5", "V6"]
+  "qwen3-4b-instruct":
+    hf_id: "Qwen/Qwen3-4B"
+    prompt:
+      start_of_turn: "<|im_start|>"
+      end_of_turn: "<|im_end|>"
+      roles:
+        user: "user"
+        model: "assistant"
+      enforce_bos: false
+      enforce_eos: false
+      allow_multiple_eos: true
+    tokenizer:
+      pad_token_strategy: "existing"
+      require_bos: false
+      require_eos: false
+      use_fast: true
+      add_prefix_space: false
+    architecture:
+      wrapper_class: "ecg_qwen_model.ECGQwenPrefix"
+      hidden_size_attrs:
+        - "config.hidden_size"
+      language_model_path_hints:
+        - "model.layers"
+        - "model.model.layers"
+      attn_implementation: "eager"
+      conv_out_dim: 64
+    lora_policy:
+      expect_language_only: true
+      allowed_markers:
+        - "model.layers."
+        - ".layers."
+      blocked_markers:
+        - "vision"
+        - "multi_modal"
+        - "projector"
+        - ".enc."
+        - "encoder_proj"
+      freeze_vision: false
+    packing:
+      prompt_format: "qwen_chat_v1"
+      conversation:
+        user_role_aliases: ["human", "user"]
+        model_role_aliases: ["gpt", "assistant"]
+        strip_image_from_roles: ["human"]
+        merge_system_with_first_user: true
+      ecg_tokens:
+        global_start: "<ecg_global_start>"
+        global_end: "<ecg_global_end>"
+        lead_start_template: "<ecg_lead_{lead_lower}_start>"
+        lead_end_template: "<ecg_lead_{lead_lower}_end>"
+        canonical_leads: ["I", "II", "III", "aVR", "aVL", "aVF", "V1", "V2", "V3", "V4", "V5", "V6"]
+  gpt-oss-120b:
+    hf_id: "openai/gpt-oss-120b"
+    prompt:
+      start_of_turn: "<|start|>"
+      end_of_turn: "<|end|>"
+      roles:
+        user: "user"
+        model: "assistant"
+      enforce_bos: false
+      enforce_eos: false
+      allow_multiple_eos: true
+      system_prompt: |-
+        You are ChatGPT, a large language model trained by OpenAI.
+        Knowledge cutoff: 2024-06
+        Current date: 2025-06-28
+
+        Reasoning: low
+
+        # Valid channels: final. Channel must be included for every message.
+      developer_prompt: |-
+        # Instructions
+
+        You are trained to interpret electrocardiograms (ECGs) and must answer questions about them clearly and accurately.
+    tokenizer:
+      pad_token_strategy: "existing"
+      require_bos: false
+      require_eos: false
+      use_fast: true
+      add_prefix_space: false
+    architecture:
+      wrapper_class: "ecg_gptoss_model.ECGGptOssPrefix"
+      hidden_size_attrs:
+        - "config.hidden_size"
+      language_model_path_hints:
+        - "model.layers"
+        - "model.model.layers"
+      attn_implementation: "eager"
+      conv_out_dim: 64
+    lora_policy:
+      expect_language_only: true
+      allowed_markers:
+        - "model.layers."
+        - ".layers."
+      blocked_markers:
+        - "vision"
+        - "multi_modal"
+        - "projector"
+        - ".enc."
+        - "encoder_proj"
+      freeze_vision: false
+    packing:
+      prompt_format: "harmony_chat_v1"
+      conversation:
+        user_role_aliases: ["human", "user"]
+        model_role_aliases: ["assistant", "gpt"]
+        strip_image_from_roles: ["human"]
+        merge_system_with_first_user: false
+      ecg_tokens:
+        global_start: "<ecg_global_start>"
+        global_end: "<ecg_global_end>"
+        lead_start_template: "<ecg_lead_{lead_lower}_start>"
+        lead_end_template: "<ecg_lead_{lead_lower}_end>"
+        canonical_leads: ["I", "II", "III", "aVR", "aVL", "aVF", "V1", "V2", "V3", "V4", "V5", "V6"]
+  gemma-1b-it:
+    hf_id: "google/gemma-3-1b-it"
+    prompt:
+      start_of_turn: "<start_of_turn>"
+      end_of_turn: "<end_of_turn>\n"
+      roles:
+        user: "user"
+        model: "model"
+      enforce_bos: true
+      enforce_eos: true
+    tokenizer:
+      pad_token_strategy: "eos"
+      require_bos: true
+      require_eos: true
+      use_fast: true
+      add_prefix_space: false
+    architecture:
+      wrapper_class: "ecg_gemma_model.ECGGemmaPrefix"
+      hidden_size_attrs:
+        - "config.hidden_size"
+        - "config.text_config.hidden_size"
+      language_model_path_hints:
+        - "base_model.model.language_model.layers"
+        - "model.language_model.layers"
+        - "model.layers"
+      attn_implementation: "eager"
+      conv_out_dim: 64
+    lora_policy:
+      expect_language_only: true
+      allowed_markers:
+        - "language_model"
+        - ".model.layers."
+        - ".layers."
+      blocked_markers:
+        - "vision"
+        - "multi_modal"
+        - "projector"
+        - ".enc."
+        - "encoder_proj"
+      freeze_vision: true
+    packing:
+      prompt_format: "gemma_chat_v1"
+      conversation:
+        user_role_aliases: ["human", "user"]
+        model_role_aliases: ["gpt", "assistant"]
+        strip_image_from_roles: ["human"]
+        merge_system_with_first_user: true
+      ecg_tokens:
+        global_start: "<ecg_global_start>"
+        global_end: "<ecg_global_end>"
+        lead_start_template: "<ecg_lead_{lead_lower}_start>"
+        lead_end_template: "<ecg_lead_{lead_lower}_end>"
+        canonical_leads: ["I", "II", "III", "aVR", "aVL", "aVF", "V1", "V2", "V3", "V4", "V5", "V6"]
+  gemma-270m-it:
+    hf_id: "google/gemma-3-270m-it"
+    prompt:
+      start_of_turn: "<start_of_turn>"
+      end_of_turn: "<end_of_turn>\n"
+      roles:
+        user: "user"
+        model: "model"
+      enforce_bos: true
+      enforce_eos: true
+    tokenizer:
+      pad_token_strategy: "eos"
+      require_bos: true
+      require_eos: true
+      use_fast: true
+      add_prefix_space: false
+    architecture:
+      wrapper_class: "ecg_gemma_model.ECGGemmaPrefix"
+      hidden_size_attrs:
+        - "config.hidden_size"
+        - "config.text_config.hidden_size"
+      language_model_path_hints:
+        - "base_model.model.language_model.layers"
+        - "model.language_model.layers"
+        - "model.layers"
+      attn_implementation: "eager"
+      conv_out_dim: 64
+    lora_policy:
+      expect_language_only: true
+      allowed_markers:
+        - "language_model"
+        - ".model.layers."
+        - ".layers."
+      blocked_markers:
+        - "vision"
+        - "multi_modal"
+        - "projector"
+        - ".enc."
+        - "encoder_proj"
+      freeze_vision: true
+    packing:
+      prompt_format: "gemma_chat_v1"
+      conversation:
+        user_role_aliases: ["human", "user"]
+        model_role_aliases: ["gpt", "assistant"]
+        strip_image_from_roles: ["human"]
+        merge_system_with_first_user: true
+      ecg_tokens:
+        global_start: "<ecg_global_start>"
+        global_end: "<ecg_global_end>"
+        lead_start_template: "<ecg_lead_{lead_lower}_start>"
+        lead_end_template: "<ecg_lead_{lead_lower}_end>"
+        canonical_leads: ["I", "II", "III", "aVR", "aVL", "aVF", "V1", "V2", "V3", "V4", "V5", "V6"]
+  medgemma-27b-it:
+    hf_id: "google/medgemma-27b-text-it"
+    prompt:
+      start_of_turn: "<start_of_turn>"
+      end_of_turn: "<end_of_turn>\n"
+      roles:
+        user: "user"
+        model: "model"
+      enforce_bos: true
+      enforce_eos: true
+    tokenizer:
+      pad_token_strategy: "eos"
+      require_bos: true
+      require_eos: true
+      use_fast: true
+      add_prefix_space: false
+    architecture:
+      wrapper_class: "ecg_gemma_model.ECGGemmaPrefix"
+      hidden_size_attrs:
+        - "config.hidden_size"
+        - "config.text_config.hidden_size"
+      language_model_path_hints:
+        - "base_model.model.language_model.layers"
+        - "model.language_model.layers"
+        - "model.layers"
+      attn_implementation: "eager"
+      conv_out_dim: 64
+    lora_policy:
+      expect_language_only: true
+      allowed_markers:
+        - "language_model"
+        - ".model.layers."
+        - ".layers."
+      blocked_markers:
+        - "vision"
+        - "multi_modal"
+        - "projector"
+        - ".enc."
+        - "encoder_proj"
+      freeze_vision: true
+    packing:
+      prompt_format: "gemma_chat_v1"
+      conversation:
+        user_role_aliases: ["human", "user"]
+        model_role_aliases: ["gpt", "assistant"]
+        strip_image_from_roles: ["human"]
+        merge_system_with_first_user: true
+      ecg_tokens:
+        global_start: "<ecg_global_start>"
+        global_end: "<ecg_global_end>"
+        lead_start_template: "<ecg_lead_{lead_lower}_start>"
+        lead_end_template: "<ecg_lead_{lead_lower}_end>"
+        canonical_leads: ["I", "II", "III", "aVR", "aVL", "aVF", "V1", "V2", "V3", "V4", "V5", "V6"]

camel_inference/src/camel/process_ecg.py

@@ -0,0 +1,208 @@
+from typing import Optional, List, Dict, Any
+import numpy as np
+import torch
+
+from read_ecg import load_record
+
+_LEAD_SYNONYMS: Dict[str, str] = {
+    # Limb
+    "I": "I", "II": "II", "III": "III",
+    "DI": "I", "DII": "II", "DIII": "III",
+    "MLII": "II",
+    # Augmented
+    "AVR": "aVR", "AVL": "aVL", "AVF": "aVF",
+    # Precordial
+    "V1": "V1", "V2": "V2", "V3": "V3", "V4": "V4", "V5": "V5", "V6": "V6",
+    # Dataset-specific
+    "ECG": "I",      # Apnea-ECG
+    "ECG1": "I", "ECG2": "II", # AFDB
+    "CM5": "V5", "D3": "V3", "D4": "V4",
+    "CM2": "V2", "ML5": "V5",
+    "VF": "VF",
+}
+
+class NormOp:
+    def __init__(self, name: str, params: Dict[str, float] | None = None):
+        self.name = name
+        self.params = params or {}
+
+def to_canonical_lead(name: str) -> Optional[str]:
+    if not isinstance(name, str):
+        return None
+    s = name.strip().upper()
+    if s in _LEAD_SYNONYMS:
+        return _LEAD_SYNONYMS[s]
+    if s in ("A VR", "A VL", "A VF"):
+        return "a" + s.replace(" ", "")[1:]
+    return None
+
+def parse_pipeline(spec: Optional[str]) -> List[NormOp]:
+    if not spec:
+        return [NormOp("nonfinite_to_zero"), NormOp("clip", {})]
+    ops: List[NormOp] = []
+    for token in spec.split(','):
+        token = token.strip()
+        if not token:
+            continue
+        parts = token.split(':')
+        name = parts[0].lower()
+        if name == 'clip':
+            mn = None; mx = None
+            if len(parts) >= 2 and parts[1] != '':
+                mn = float(parts[1])
+            if len(parts) >= 3 and parts[2] != '':
+                mx = float(parts[2])
+            ops.append(NormOp('clip', {'min': mn, 'max': mx}))
+        elif name == 'nonfinite_to_zero':
+            ops.append(NormOp('nonfinite_to_zero'))
+        else:
+            raise ValueError(f"Unknown op '{name}'.")
+    return ops
+
+def apply_pre_ops(signal_1d: np.ndarray, ops: List[NormOp]) -> np.ndarray:
+    x = signal_1d
+    for op in ops:
+        if op.name == 'nonfinite_to_zero':
+            np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0, copy=False)
+    return x
+
+
+def apply_post_ops(segments: np.ndarray, ops: List[NormOp], lead_name: str,
+                   clip_stats: Optional[Dict[str, Dict[str, float]]] = None) -> np.ndarray:
+    x = segments
+    for op in ops:
+        if op.name == 'clip':
+            mn = op.params.get('min', None)
+            mx = op.params.get('max', None)
+            if mn is None and mx is None and clip_stats is not None and lead_name in clip_stats:
+                mn = clip_stats[lead_name].get('clip_min', None)
+                mx = clip_stats[lead_name].get('clip_max', None)
+            if mn is None:
+                mn = -np.inf
+            if mx is None:
+                mx = np.inf
+            if mn > mx:
+                mn, mx = mx, mn
+            np.clip(x, mn, mx, out=x)
+    return x.astype(np.float32, copy=False)
+
+def _segment_data(ecg_signal: np.ndarray, raw_fs: int) -> np.ndarray:
+    """
+    Segment 1D signal sampled at ``raw_fs`` into 1-second clips and resample to 256 samples.
+    Returns np.float32 array [N_segments, 256].
+    """
+    assert raw_fs > 0, f"raw sampling rate must be positive (got {raw_fs})"
+    samples_per_second = int(raw_fs)
+    n_samples = int(len(ecg_signal))
+    if samples_per_second <= 0 or n_samples <= 0:
+        return np.empty((0, 256), dtype=np.float32)
+
+    # Pre-allocate upper bound
+    n_full = n_samples // samples_per_second
+    has_partial = (n_samples % samples_per_second) > 0
+    max_segments = n_full + (1 if has_partial else 0)
+    if max_segments == 0:
+        return np.empty((0, 256), dtype=np.float32)
+
+    out = np.empty((max_segments, 256), dtype=np.float32)
+    new_idx = np.linspace(0, 1, num=256, dtype=np.float32)
+    actual = 0
+    for start in range(0, n_samples, samples_per_second):
+        end = min(start + samples_per_second, n_samples)
+        seg = ecg_signal[start:end]
+        if seg.shape[0] < samples_per_second * 0.5:
+            continue
+        old_idx = np.linspace(0, 1, num=seg.shape[0], dtype=np.float32)
+        out[actual] = np.interp(new_idx, old_idx, seg).astype(np.float32, copy=False)
+        actual += 1
+    return out[:actual]
+
+def _apply_filters(signal_1d, fs: int):
+    """
+    Apply ECG signal filters: 50/60 Hz notch filters and 0.3 Hz high-pass filter.
+    
+    Removes powerline interference and baseline wander from ECG signals.
+    Uses cascaded second-order sections (SOS) for numerical stability.
+    """
+    import numpy as np
+    from scipy.signal import iirnotch, butter, sosfiltfilt, tf2sos
+
+    x = np.asarray(signal_1d, dtype=np.float64)
+    if x.size == 0 or fs <= 0:
+        return x.astype('float32')
+
+    # Design notch filters for powerline interference
+    Q = 30.0
+    nyq = fs / 2.0
+    sos_filters = []
+    for freq in (50.0, 60.0):
+        if freq >= nyq:
+            continue
+        b, a = iirnotch(freq, Q, fs)
+        sos_filters.append(tf2sos(b, a))
+    
+    # Add high-pass filter for baseline wander removal
+    sos_hp = butter(N=2, Wn=0.3 / nyq, btype='highpass', output='sos')
+    sos_filters.append(sos_hp)
+    
+    # Apply all filters in cascade
+    sos = np.vstack(sos_filters)
+    x = sosfiltfilt(sos, x)
+    return x.astype('float32')
+
+def get_waveform(device:torch.device, ecg_path:str, start_sec = None, end_sec = None, leads: Optional[List[str]] = None,
+                 process:bool = False, norm: str = "nonfinite_to_zero,clip") -> Dict[str, Any]:
+    """
+    Run ECG preprocessing pipeline (Step II): raw → filter → segment → normalize → storage.
+    
+    Processes ECG records from manifest, applies signal filtering, segments into 1s windows,
+    normalizes using computed clip stats, and writes to efficient storage format.
+    
+    Args:
+        dataset: Dataset name
+        p: file path
+        clip_stats_path: Path to clip_stats json file
+        process: Y/N process
+        fs: Target sampling frequency
+        leads: Optional list of leads to process
+        norm: Normalization pipeline (e.g., "nonfinite_to_zero,clip")
+    
+    Returns:
+        Dict with output paths: {output_dir, index, clip_stats}
+    """
+
+    # Parallel path: use per-worker shard/npy directories then aggregate
+    ops = parse_pipeline(norm)
+
+    # Convert to dict format
+    ecg_dict = {}
+    try:
+        df, sig_names, original_fs = load_record(ecg_path, start_sec, end_sec, leads)
+        
+        # Process each lead in the record
+        for i, raw_name in enumerate(sig_names):
+            canon = to_canonical_lead(raw_name)
+            if not canon:
+                print('to_canonical_lead')
+                continue
+            x = df[:, i].astype('float32', copy=False)
+
+            if process:
+                x = apply_pre_ops(x, ops)  # Handle non-finite values
+                x = _apply_filters(x, original_fs)  # Remove noise and baseline wander
+
+            # Segment into 1s windows and normalize
+            segs = _segment_data(x, original_fs)
+            segs = apply_post_ops(segs, ops, canon)
+
+            lead_tensor = torch.from_numpy(segs).to(torch.float32).to(device)
+            if not (torch.any(lead_tensor.isnan())):
+                lead_tensor = lead_tensor.nan_to_num()
+
+            ecg_dict[canon] = lead_tensor
+    
+    except Exception as e:
+        # Skip failed records
+        print(e)
+
+    return ecg_dict

camel_inference/src/camel/projectors.py

@@ -0,0 +1,57 @@
+"""Projector registry utilities.
+
+Provides a lightweight mechanism to swap the adapter architecture that maps the
+conv encoder output to the language-model hidden size. Mirrors the ergonomic
+API used by loss.py: a registry, default implementations, and a simple factory.
+"""
+from __future__ import annotations
+
+from typing import Callable, Dict, Iterable
+import torch.nn as nn
+
+ProjectorBuilder = Callable[[int, int], nn.Module]
+
+_PROJECTOR_REGISTRY: Dict[str, ProjectorBuilder] = {}
+
+def register_projector(name: str) -> Callable[[ProjectorBuilder], ProjectorBuilder]:
+    """Decorator to register a projector builder under a unique name."""
+    key = name.strip().lower()
+
+    def _decorator(fn: ProjectorBuilder) -> ProjectorBuilder:
+        if not callable(fn):
+            raise TypeError("Projector builder must be callable.")
+        if key in _PROJECTOR_REGISTRY:
+            raise ValueError(f"Projector '{name}' is already registered.")
+        _PROJECTOR_REGISTRY[key] = fn
+        return fn
+
+    return _decorator
+
+@register_projector("linear")
+def _linear_projector(in_dim: int, out_dim: int) -> nn.Module:
+    """Single linear adapter (current default)."""
+    return nn.Linear(in_dim, out_dim, bias=True)
+
+def available_projectors() -> Iterable[str]:
+    """Return sorted projector names."""
+    return sorted(_PROJECTOR_REGISTRY.keys())
+
+def build_projector(name: str, in_dim: int, out_dim: int) -> nn.Module:
+    """Instantiate a registered projector."""
+    if not _PROJECTOR_REGISTRY:
+        raise RuntimeError("No projectors registered.")
+    key = (name or "").strip().lower()
+    if not key:
+        raise ValueError("Projector name must be a non-empty string.")
+    builder = _PROJECTOR_REGISTRY.get(key)
+    if builder is None:
+        raise KeyError(
+            f"Unknown projector '{name}'. Available: {', '.join(available_projectors())}"
+        )
+    return builder(int(in_dim), int(out_dim))
+
+__all__ = [
+    "ProjectorBuilder",
+    "available_projectors",
+    "build_projector",
+]

camel_inference/src/camel/prompt_renderers.py

@@ -0,0 +1,209 @@
+"""Prompt rendering and span construction helpers."""
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, List, Tuple
+from transformers import PreTrainedTokenizer
+
+from camel.assertions import (
+    assert_tokenization_cursor_matches,
+    assert_model_spans_valid,
+    assert_eos_appended,
+)
+
+def _ensure_trailing_newline(s: str) -> str:
+    if s.endswith("\n"):
+        return s
+    return s + "\n"
+
+def _chat_v1_wrappers(schema, role: str) -> Tuple[str, str]:
+    prefix = f"{schema.prompt.start_of_turn}{role}\n"
+    suffix = _ensure_trailing_newline(schema.prompt.end_of_turn)
+    return prefix, suffix
+
+def _harmony_v1_wrappers(schema, role: str, *, use_return: bool = False) -> Tuple[str, str]:
+    if role == schema.prompt.model_role:
+        prefix = f"{schema.prompt.start_of_turn}{role}"
+        suffix = "<|return|>" if use_return else str(schema.prompt.end_of_turn)
+        return prefix, suffix
+    prefix = f"{schema.prompt.start_of_turn}{role}<|message|>"
+    suffix = str(schema.prompt.end_of_turn)
+    return prefix, suffix
+
+def _render_with_wrappers(
+    tokenizer: PreTrainedTokenizer,
+    turns: List[Dict[str, str]],
+    *,
+    schema,
+    wrapper_fn,
+) -> Dict[str, Any]:
+    tok = tokenizer
+    prompt_tokens = schema.prompt
+
+    text_ids: List[int] = []
+    if prompt_tokens.require_bos and tok.bos_token_id is not None:
+        text_ids.append(tok.bos_token_id)
+
+    model_spans_in_text: List[Tuple[int, int]] = []
+    cursor = len(text_ids)
+    text_preview_parts: List[str] = []
+
+    for turn in turns:
+        role = turn["role"]
+        text_block = turn["text_block"]
+        prefix, suffix = wrapper_fn(schema, role)
+        content = text_block
+        if suffix and content.endswith(suffix):
+            content = content[: -len(suffix)]
+        ids_prefix = tok.encode(prefix, add_special_tokens=False)
+        ids_content = tok.encode(content, add_special_tokens=False)
+        ids_suffix = tok.encode(suffix, add_special_tokens=False)
+
+        text_ids.extend(ids_prefix)
+        text_ids.extend(ids_content)
+        text_ids.extend(ids_suffix)
+
+        if role == prompt_tokens.model_role:
+            s = cursor + len(ids_prefix)
+            e = s + len(ids_content) + len(ids_suffix)
+            if e > s:
+                model_spans_in_text.append((s, e))
+        cursor += len(ids_prefix) + len(ids_content) + len(ids_suffix)
+        text_preview_parts.append(prefix + content + suffix)
+
+    assert_tokenization_cursor_matches(cursor, len(text_ids))
+
+    if prompt_tokens.require_eos and tok.eos_token_id is not None:
+        text_ids.append(tok.eos_token_id)
+        if model_spans_in_text and turns[-1]["role"] == prompt_tokens.model_role:
+            model_spans_in_text[-1] = (model_spans_in_text[-1][0], len(text_ids))
+
+    assert_eos_appended(text_ids, tok, prompt_tokens.require_eos)
+    assert_model_spans_valid(model_spans_in_text, len(text_ids))
+
+    return {
+        "text_ids": text_ids,
+        "model_spans_in_text": model_spans_in_text,
+        "text_preview": "".join(text_preview_parts),
+    }
+
+def _render_chat_v1(
+    tokenizer: PreTrainedTokenizer,
+    turns: List[Dict[str, str]],
+    *,
+    schema,
+) -> Dict[str, Any]:
+    return _render_with_wrappers(
+        tokenizer,
+        turns,
+        schema=schema,
+        wrapper_fn=_chat_v1_wrappers,
+    )
+
+def _render_harmony_v1(
+    tokenizer: PreTrainedTokenizer,
+    turns: List[Dict[str, str]],
+    *,
+    schema,
+    use_return_for_last_assistant: bool = False,
+) -> Dict[str, Any]:
+    tok = tokenizer
+    prompt_tokens = schema.prompt
+
+    text_ids: List[int] = []
+    if prompt_tokens.require_bos and tok.bos_token_id is not None:
+        text_ids.append(tok.bos_token_id)
+
+    model_spans_in_text: List[Tuple[int, int]] = []
+    cursor = len(text_ids)
+    text_preview_parts: List[str] = []
+
+    last_assistant_idx = None
+    if use_return_for_last_assistant:
+        for idx in range(len(turns) - 1, -1, -1):
+            if turns[idx]["role"] == prompt_tokens.model_role:
+                last_assistant_idx = idx
+                break
+
+    for idx, turn in enumerate(turns):
+        role = turn["role"]
+        text_block = turn["text_block"]
+        use_return = use_return_for_last_assistant and last_assistant_idx is not None and idx == last_assistant_idx
+        prefix, suffix = _harmony_v1_wrappers(schema, role, use_return=use_return)
+        content = text_block
+        if suffix and content.endswith(suffix):
+            content = content[: -len(suffix)]
+        ids_prefix = tok.encode(prefix, add_special_tokens=False)
+        ids_content = tok.encode(content, add_special_tokens=False)
+        ids_suffix = tok.encode(suffix, add_special_tokens=False)
+
+        text_ids.extend(ids_prefix)
+        text_ids.extend(ids_content)
+        text_ids.extend(ids_suffix)
+
+        if role == prompt_tokens.model_role:
+            s = cursor + len(ids_prefix)
+            e = s + len(ids_content) + len(ids_suffix)
+            if e > s:
+                model_spans_in_text.append((s, e))
+        cursor += len(ids_prefix) + len(ids_content) + len(ids_suffix)
+        text_preview_parts.append(prefix + content + suffix)
+
+    assert_tokenization_cursor_matches(cursor, len(text_ids))
+
+    if prompt_tokens.require_eos and tok.eos_token_id is not None:
+        text_ids.append(tok.eos_token_id)
+        if model_spans_in_text and turns[-1]["role"] == prompt_tokens.model_role:
+            model_spans_in_text[-1] = (model_spans_in_text[-1][0], len(text_ids))
+
+    assert_eos_appended(text_ids, tok, prompt_tokens.require_eos)
+    assert_model_spans_valid(model_spans_in_text, len(text_ids))
+
+    return {
+        "text_ids": text_ids,
+        "model_spans_in_text": model_spans_in_text,
+        "text_preview": "".join(text_preview_parts),
+    }
+
+_PROMPT_RENDERERS: Dict[str, Callable[[PreTrainedTokenizer, List[Dict[str, str]], Any], Dict[str, Any]]] = {
+    "gemma_chat_v1": _render_chat_v1,
+    "qwen_chat_v1": _render_chat_v1,
+}
+
+def render_prompt_and_spans(
+    tokenizer: PreTrainedTokenizer,
+    turns: List[Dict[str, str]],
+    *,
+    schema,
+    use_return_for_last_assistant: bool = False,
+) -> Dict[str, Any]:
+    format_id = str(schema.conversation.format_id)
+    if format_id == "harmony_chat_v1":
+        return _render_harmony_v1(
+            tokenizer,
+            turns,
+            schema=schema,
+            use_return_for_last_assistant=use_return_for_last_assistant,
+        )
+    renderer = _PROMPT_RENDERERS.get(format_id)
+    if renderer is None:
+        raise ValueError(f"Unknown prompt format '{format_id}'.")
+    return renderer(tokenizer, turns, schema=schema)
+
+def turn_wrappers(schema, role: str, *, use_return: bool = False) -> Tuple[str, str]:
+    format_id = str(schema.conversation.format_id)
+    if format_id in ("gemma_chat_v1", "qwen_chat_v1"):
+        return _chat_v1_wrappers(schema, role)
+    if format_id == "harmony_chat_v1":
+        return _harmony_v1_wrappers(schema, role, use_return=use_return)
+    raise ValueError(f"Unknown prompt format '{format_id}'.")
+
+def assistant_generation_prefix(schema) -> str:
+    format_id = str(schema.conversation.format_id)
+    if format_id in ("gemma_chat_v1", "qwen_chat_v1"):
+        return f"{schema.prompt.start_of_turn}{schema.prompt.model_role}\n"
+    if format_id == "harmony_chat_v1":
+        return f"{schema.prompt.start_of_turn}{schema.prompt.model_role}"
+    raise ValueError(f"Unknown prompt format '{format_id}'.")
+
+
+__all__ = ["render_prompt_and_spans", "turn_wrappers", "assistant_generation_prefix"]

camel_inference/src/camel/training_setup.py

@@ -0,0 +1,140 @@
+"""
+Runtime and configuration helpers extracted from train_ecg_text.py.
+These utilities keep the training entrypoint concise while preserving the
+original behaviour when preparing distributed state, tokenizer metadata, and
+packing configuration.
+"""
+from __future__ import annotations
+
+from typing import Dict, Optional, List
+import torch.distributed as dist
+from transformers import AutoTokenizer
+
+from camel.ecg_text_packing import (
+    ECGSpecialTokenCatalog,
+    PackingSchema,
+    PromptTokens,
+)
+from camel.model_registry import ModelConfig, ModelRegistryError, load_registry
+
+def is_main_process() -> bool:
+    """Return True for rank 0 (or standalone execution)."""
+    return (not dist.is_initialized()) or dist.get_rank() == 0
+
+def build_packing_schema(pretrained_model_id: str) -> PackingSchema:
+    """
+    Construct the packing schema (prompt + conversation rules + ECG tokens)
+    for the given backbone using the shared registry.
+    """
+    registry = load_registry()
+    cfg: Optional[ModelConfig]
+    try:
+        cfg = registry.get(pretrained_model_id)
+    except ModelRegistryError:
+        cfg = None
+        for name in registry.names():
+            candidate = registry.get(name)
+            if candidate.hf_id == pretrained_model_id:
+                cfg = candidate
+                break
+    if cfg is None:
+        raise ModelRegistryError(
+            f"Pretrained model '{pretrained_model_id}' not found in registry at {registry.source_path}"
+        )
+    prompt_cfg = cfg.prompt_config()
+    roles_dict = dict(prompt_cfg.roles or {})
+    try:
+        user_role = str(roles_dict["user"])
+        model_role = str(roles_dict["model"])
+    except KeyError as exc:
+        missing = exc.args[0]
+        raise ModelRegistryError(
+            f"Prompt configuration for registry entry '{cfg.name}' is missing the '{missing}' role."
+        ) from exc
+    prompt_tokens = PromptTokens(
+        start_of_turn=prompt_cfg.start_of_turn,
+        end_of_turn=prompt_cfg.end_of_turn,
+        user_role=user_role,
+        model_role=model_role,
+        require_bos=prompt_cfg.enforce_bos,
+        require_eos=prompt_cfg.enforce_eos,
+        allow_multiple_eos=prompt_cfg.allow_multiple_eos,
+    )
+    packing_cfg = cfg.packing_config()
+    conversation_rules = packing_cfg.conversation
+    ecg_tokens = packing_cfg.ecg_tokens
+    return PackingSchema(
+        prompt=prompt_tokens,
+        conversation=conversation_rules,
+        ecg=ecg_tokens,
+    )
+
+def initialize_tokenizer(
+    model_id: str,
+    *,
+    trust_remote_code: bool = True,
+    use_fast: Optional[bool] = None,
+    add_prefix_space: Optional[bool] = None,
+) -> AutoTokenizer:
+    """
+    Instantiate the HF tokenizer, allowing policy to be driven by the registry
+    (use_fast/add_prefix_space). If not provided, defaults are use_fast=True,
+    add_prefix_space=False.
+    """
+    # Honor registry defaults when the caller doesn't override them.
+    default_use_fast = True
+    default_add_prefix_space = False
+    try:
+        registry = load_registry()
+        cfg: Optional[ModelConfig]
+        try:
+            cfg = registry.get(model_id)
+        except ModelRegistryError:
+            cfg = None
+            for name in registry.names():
+                candidate = registry.get(name)
+                if candidate.hf_id == model_id:
+                    cfg = candidate
+                    break
+        if cfg is not None:
+            tcfg = cfg.tokenizer_config()
+            default_use_fast = bool(tcfg.use_fast)
+            default_add_prefix_space = bool(tcfg.add_prefix_space)
+    except Exception:
+        # Fall back to built-in defaults if registry is unavailable.
+        pass
+
+    return AutoTokenizer.from_pretrained(
+        model_id,
+        use_fast=default_use_fast if use_fast is None else bool(use_fast),
+        add_prefix_space=default_add_prefix_space if add_prefix_space is None else bool(add_prefix_space),
+        trust_remote_code=trust_remote_code,
+    )
+
+def register_ecg_special_tokens(
+    tokenizer: AutoTokenizer,
+    catalog: ECGSpecialTokenCatalog,
+) -> Dict[int, int]:
+    """
+    Ensure the tokenizer includes the ECG special tokens from the provided catalog.
+    Returns a mapping from catalog index to token ID.
+    """
+    # Add only tokens that are currently unknown to the tokenizer (not present
+    # as core specials or regular vocab entries).
+    tokens_to_add: List[str] = []
+    for token in catalog.tokens:
+        tok_id = tokenizer.convert_tokens_to_ids(token)
+        if tok_id is None or tok_id == tokenizer.unk_token_id:
+            tokens_to_add.append(token)
+    if tokens_to_add:
+        tokenizer.add_special_tokens({"additional_special_tokens": tokens_to_add})
+    ecg_special_token_id_map: Dict[int, int] = {}
+    for token, catalog_index in catalog.token_to_index.items():
+        token_id = tokenizer.convert_tokens_to_ids(token)
+        if token_id is None or token_id == tokenizer.unk_token_id:
+            raise RuntimeError(f"Tokenizer failed to register ECG special token: {token}")
+        encoded = tokenizer.encode(token, add_special_tokens=False)
+        if len(encoded) != 1 or encoded[0] != token_id:
+            raise RuntimeError(f"ECG special token does not map to a single id: {token}")
+        ecg_special_token_id_map[catalog_index] = int(token_id)
+    return ecg_special_token_id_map

camel_inference/src/read_ecg.py

@@ -0,0 +1,33 @@
+from typing import Optional
+import numpy as np
+import wfdb
+
+def load_record(ecg_path, start_sec: Optional[int], end_sec: Optional[int], leads: Optional[list[str]]):
+    record = wfdb.rdrecord(ecg_path)
+    fs = record.fs
+    lead_names = record.sig_name
+
+    signal = record.p_signal # n_samples x n_leads
+    if leads:
+        kept_signals, kept_leads = [], []
+        lead_to_idx = {name: i for i, name in enumerate(lead_names)}
+        for l in leads:
+            if l in lead_to_idx:
+                kept_signals.append(signal[:, lead_to_idx[l]])
+                kept_leads.append(l)
+            else:
+                print(f'Lead {l} does not exist. Skipping.')
+        if not kept_signals:
+            raise ValueError(f"None of the requested leads were found. requested={leads}, available={lead_names}")
+        
+        signal = np.stack(kept_signals, axis=1)
+        lead_names = kept_leads
+
+    # Optinally subsample the signal
+    start_ind = 0 if start_sec is None else start_sec * fs
+    end_ind = len(signal) if end_sec is None else end_sec * fs
+    if end_ind > len(signal):
+        print(f'ECG is {len(signal) / fs} seconds')
+    signal = signal[start_ind:end_ind, :]
+
+    return signal, lead_names, fs