Bundle camel library, demo data, and update app.py for HF Space

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 import os
 import sys
-import subprocess
 from pathlib import Path
 import wfdb
 import numpy as np
@@ -9,308 +8,117 @@ import matplotlib.pyplot as plt
 import torch
 import glob
 from types import SimpleNamespace
+from huggingface_hub import hf_hub_download
 
-# --- Configuration ---
-REPO_URL = "https://github.com/CAMEL-ECG/CAMEL-inference.git"
-REPO_DIR = Path("./camel_inference")
-DEMO_DIR = REPO_DIR / "demo"
+# --- Setup: Checkpoints & Model ---
+CHECKPOINT_DIR = Path("checkpoints")
+CHECKPOINT_DIR.mkdir(exist_ok=True)
 
-# --- Setup: Clone Inference Repo & Install ---
-def setup_inference_env():
-    """Clones the repo and ensures it's in sys.path."""
-    if not REPO_DIR.exists():
-        print(f"Cloning {REPO_URL}...")
-        subprocess.run(["git", "clone", REPO_URL, str(REPO_DIR)], check=True)
-    else:
-        if (REPO_DIR / ".git").exists():
-            print("Repo directory exists. Pulling latest...")
-            try:
-                subprocess.run(["git", "-C", str(REPO_DIR), "pull"], check=True)
-            except Exception as e:
-                print(f"Git pull failed (non-fatal): {e}")
-        else:
-            print("Found bundled inference code. Skipping git pull.")
+# Ensure checkpoints are symlinked from cache to local dir for CAMEL
+files = ["camel_base.pt", "camel_ecginstruct.pt", "camel_forecast.pt"]
+repo_id = "CAMEL-ECG/CAMEL"
 
-    # Ensure the src is importable
-    if str(REPO_DIR.resolve()) not in sys.path:
-        sys.path.append(str(REPO_DIR.resolve()))
-    
-    # Install the package in editable mode if not already (for dependencies)
-    # Note: In a persistent space, we might check if it's already installed.
-    # For now, we rely on requirements.txt satisfying most, but we run this to be safe.
-    try:
-        subprocess.run([sys.executable, "-m", "pip", "install", "-e", str(REPO_DIR)], check=True)
-    except Exception as e:
-        print(f"Pip install failed: {e}")
+for f in files:
+    cached_path = hf_hub_download(repo_id=repo_id, filename=f)
+    target_link = CHECKPOINT_DIR / f
+    if not target_link.exists():
+        target_link.symlink_to(cached_path)
+        print(f"Symlinked {f} -> {cached_path}")
 
-def download_checkpoints():
-    """Downloads CAMEL checkpoints if missing."""
-    from huggingface_hub import hf_hub_download
-    import shutil
-    
-    checkpoint_dir = Path("checkpoints")
-    checkpoint_dir.mkdir(exist_ok=True)
-    
-    repo_id = "CAMEL-ECG/CAMEL"
-    files = ["camel_base.pt", "camel_ecginstruct.pt", "camel_forecast.pt"]
-    
-    for f in files:
-        dest = checkpoint_dir / f
-        if not dest.exists():
-            print(f"Downloading {f}...")
-            try:
-                src = hf_hub_download(repo_id=repo_id, filename=f)
-                shutil.copy(src, dest)
-                print(f"Downloaded {f}")
-            except Exception as e:
-                print(f"Error downloading {f}: {e}")
-        else:
-             print(f"Checkpoint {f} exists.")
-
-setup_inference_env()
-try:
-    download_checkpoints()
-except Exception as e:
-    print(f"Checkpoint download warning: {e}")
-
-# --- Model Loading (Global Cache) ---
-# Import after setup to ensure paths are correct
-try:
-    from camel.camel_model import CAMEL
-except ImportError:
-    # Fallback if pip install -e didn't work immediately in this process
-    sys.path.append(str(REPO_DIR / "src"))
-    from camel.camel_model import CAMEL
+# Import local CAMEL library
+from camel.camel_model import CAMEL
 
 MODEL_CACHE = None
 
 def get_model():
     global MODEL_CACHE
     if MODEL_CACHE is None:
-        print("Loading CAMEL model...")
-        
-        # Verify CUDA as requested
-        print(f"Is CUDA available: {torch.cuda.is_available()}")
-        if torch.cuda.is_available():
-            print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
-            
         device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        print(f"Using device: {device}")
-        
-        # Initialize model - checkpoints will be downloaded automatically by transformers/huggingface_hub
+        print(f"Loading model on {device}...")
         MODEL_CACHE = CAMEL(mode='base', device=device)
-        print("Model loaded.")
     return MODEL_CACHE
 
 # --- Helper Functions ---
+DEMO_DIR = Path("demo")
 
 def load_examples():
-    """Returns a list of example basenames from the demo directory."""
-    if not DEMO_DIR.exists():
-        return []
-    # Search for .hea files
-    hea_files = glob.glob(str(DEMO_DIR / "*.hea"))
-    # Extract basenames (filename without extension)
-    basenames = [Path(f).stem for f in hea_files]
-    return sorted(basenames)
+    return sorted([p.stem for p in DEMO_DIR.glob("*.hea")]) if DEMO_DIR.exists() else []
 
-def read_ecg_file(file_path):
-    """Reads an ECG record using wfdb."""
-    # wfdb.rdrecord expects the path WITHOUT extension for the 'record_name' argument
-    # but we usually pass a path like "demo/12345". 
-    record_path = str(Path(file_path).with_suffix(""))
+def read_ecg(path_str):
     try:
-        record = wfdb.rdrecord(record_path)
-        return record
+        record_path = str(Path(path_str).with_suffix(""))
+        return wfdb.rdrecord(record_path)
     except Exception as e:
-        print(f"Error reading record {record_path}: {e}")
+        print(f"Error reading {path_str}: {e}")
         return None
 
 def plot_ecg(record):
-    """Plots the ECG signal."""
-    if record is None:
-        return None
-    
-    signals = record.p_signal
-    # Plot only the first lead for clarity, or a few leads
+    if not record: return None
     fig, ax = plt.subplots(figsize=(10, 4))
-    
-    # If multiple channels, plot the first one (usually I or II)
-    # Determine number of samples to plot (e.g., first 10 seconds aka 10*fs)
-    fs = record.fs
-    duration_secs = 10
-    samples = int(min(signals.shape[0], fs * duration_secs))
-    
-    t = np.arange(samples) / fs
-    ax.plot(t, signals[:samples, 0], label=record.sig_name[0])
-    
-    ax.set_title(f"ECG Signal: {record.record_name} (Lead {record.sig_name[0]})")
-    ax.set_xlabel("Time (s)")
-    ax.set_ylabel(record.units[0] if record.units else "mV")
-    ax.legend()
-    ax.grid(True, alpha=0.3)
-    
+    signals = record.p_signal
+    samples = int(min(signals.shape[0], record.fs * 10))
+    t = np.arange(samples) / record.fs
+    ax.plot(t, signals[:samples, 0])
+    ax.set_title(f"Lead {record.sig_name[0]}")
     plt.tight_layout()
     return fig
 
-def get_header_info(record):
-    """Extracts ground truth/comments from the header."""
-    if record is None:
-        return "Error reading record."
-    
-    info = []
-    if record.comments:
-        info.append("Comments/Labels:")
-        for c in record.comments:
-            info.append(f"- {c}")
-    
-    info.append(f"Fs: {record.fs} Hz")
-    info.append(f"Length: {record.sig_len} samples")
-    
-    return "\n".join(info)
+def get_info(record):
+    if not record: return ""
+    return f"Fs: {record.fs} Hz\nLength: {record.sig_len}\nComments: {record.comments}"
 
 def process_upload(files):
-    """Handles uploaded .dat/.hea files."""
-    if not files:
-        return None, "No files uploaded.", None
-
-    # We need to find the pair (.dat and .hea)
-    # Gradio uploads to a temp dir. We need to identify the common basename.
-    # Group by basename
+    if not files: return None, "", None
     paths = [Path(f.name) for f in files]
-    # Simple heuristic: take the first .hea found and look for its .dat sibling
-    hea_path = next((p for p in paths if p.suffix == '.hea'), None)
-    
-    if not hea_path:
-        return None, "Upload must include a .hea file.", None
+    hea = next((p for p in paths if p.suffix == '.hea'), None)
+    if not hea: return None, "Missing .hea file", None
     
-    # wfdb needs the path without extension
-    record = read_ecg_file(str(hea_path))
-    
-    if record is None:
-         return None, "Failed to read uploaded ECG.", None
-         
-    fig = plot_ecg(record)
-    info = get_header_info(record)
-    
-    return fig, info, str(hea_path.with_suffix("")) # Return path for inference
+    record = read_ecg(str(hea))
+    return plot_ecg(record), get_info(record), str(hea.with_suffix(""))
 
-def process_example(example_name):
-    """Handles selected example."""
-    if not example_name:
-        return None, "No example selected.", None
-        
-    path = DEMO_DIR / example_name
-    record = read_ecg_file(str(path))
-    
-    fig = plot_ecg(record)
-    info = get_header_info(record)
-    
-    return fig, info, str(path) # Return path for inference
+def process_example(name):
+    if not name: return None, "", None
+    path = DEMO_DIR / name
+    record = read_ecg(str(path))
+    return plot_ecg(record), get_info(record), str(path)
 
-def run_inference_logic(ecg_path_str, prompt="Describe the ECG."):
-    """Runs the CAMEL inference."""
-    if not ecg_path_str:
-        return "Please select an ECG first."
-        
+def run_inference(ecg_path, prompt):
+    if not ecg_path: return "Select ECG first."
     model = get_model()
-    
-    # Construct args expected by model.run()
-    # Looking at run_camel.py: 
-    # args.mode, args.text, args.ecgs (list), args.device, args.ecg_configs
-    # Defaults in run_camel.py: mode='base', device='cuda:0', top_k=64, top_p=0.95...
-    
     args = SimpleNamespace(
-        mode='base',
-        text=prompt,
-        ecgs=[ecg_path_str],
-        device=model.device, # Use model's device
-        ecg_configs=None, # Default
-        json=None,
-        temperature=0.0,
-        top_k=64,
-        top_p=0.95,
-        min_p=0.0,
-        max_new_tokens=512
+        mode='base', text=prompt, ecgs=[ecg_path], device=model.device,
+        ecg_configs=None, json=None, temperature=0.0, top_k=64, 
+        top_p=0.95, min_p=0.0, max_new_tokens=512
     )
-    
-    try:
-        output, used_prompt = model.run(args)
-        return output
-    except Exception as e:
-        return f"Inference failed: {e}"
+    output, _ = model.run(args)
+    return output
 
-
-# --- UI Construction ---
-with gr.Blocks(title="CAMEL ECG Analysis") as demo:
-    gr.Markdown("# 🐪 CAMEL ECG Foundation Model")
-    gr.Markdown("Upload an ECG (.dat and .hea) or select an example to view the signal and run inference.")
-    
-    current_ecg_path = gr.State()
+# --- UI ---
+with gr.Blocks(title="CAMEL ECG") as demo:
+    gr.Markdown("# 🐪 CAMEL ECG Model")
+    cur_path = gr.State()
     
     with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("### Input")
+        with gr.Column():
+            with gr.Tabs():
+                with gr.Tab("Example"):
+                    ex_dd = gr.Dropdown(load_examples(), label="Example")
+                    ex_btn = gr.Button("Load")
+                with gr.Tab("Upload"):
+                    up_file = gr.File(file_count="multiple")
+                    up_btn = gr.Button("Load Files")
             
-            tab_example, tab_upload = gr.Tabs(), gr.Tabs()
+            prompt = gr.Textbox("Describe the ECG.", label="Prompt")
+            run_btn = gr.Button("Run", variant="primary")
             
-            with gr.Tab("Examples"):
-                example_dropdown = gr.Dropdown(
-                    label="Choose an Example",
-                    choices=load_examples(),
-                    value=None,
-                    interactive=True
-                )
-                load_example_btn = gr.Button("Load Example", size="sm")
-            
-            with gr.Tab("Upload"):
-                file_upload = gr.File(
-                    label="Upload .dat AND .hea",
-                    file_count="multiple",
-                    file_types=[".dat", ".hea"]
-                )
-                load_upload_btn = gr.Button("Load Uploaded Files", size="sm")
-            
-            gr.Markdown("---")
-            gr.Markdown("### Inference Config")
-            prompt_input = gr.Textbox(
-                label="Prompt", 
-                value="Describe the ECG.",
-                lines=2
-            )
-            run_btn = gr.Button("Run Inference", variant="primary")
-
-        with gr.Column(scale=2):
-            gr.Markdown("### Visualization & Ground Truth")
-            plot_output = gr.Plot(label="ECG Trace")
-            label_output = gr.Code(label="Header Info / Ground Truth", language="markdown")
-            
-            gr.Markdown("### Model Output")
-            inference_output = gr.Textbox(label="CAMEL Response", interactive=False, lines=5)
-
-    # --- Interaction Logic ---
-    
-    def on_load_example(evt):
-        fig, info, path = process_example(evt)
-        return fig, info, path
+        with gr.Column():
+            plot = gr.Plot()
+            info = gr.Code(language="markdown")
+            out = gr.Textbox(label="Output", lines=5)
 
-    load_example_btn.click(
-        on_load_example,
-        inputs=[example_dropdown],
-        outputs=[plot_output, label_output, current_ecg_path]
-    )
-
-    load_upload_btn.click(
-        process_upload,
-        inputs=[file_upload],
-        outputs=[plot_output, label_output, current_ecg_path]
-    )
-    
-    run_btn.click(
-        run_inference_logic,
-        inputs=[current_ecg_path, prompt_input],
-        outputs=[inference_output]
-    )
+    ex_btn.click(process_example, ex_dd, [plot, info, cur_path])
+    up_btn.click(process_upload, up_file, [plot, info, cur_path])
+    run_btn.click(run_inference, [cur_path, prompt], out)
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
+    demo.launch(server_name="0.0.0.0", server_port=7860)

camel_inference/.gitignore

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

camel_inference/README.md

@@ -1,82 +0,0 @@
-# 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/pyproject.toml

@@ -1,36 +0,0 @@
-[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

@@ -1,49 +0,0 @@
-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

@@ -1,37 +0,0 @@
-#!/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."

requirements.txt

@@ -12,3 +12,4 @@ sentencepiece
 torch
 transformers
 wfdb
+accelerate

run_local.sh

@@ -1,6 +0,0 @@
-#!/bin/bash
-echo "Installing dependencies..."
-pip install -r requirements.txt
-
-echo "Running CAMEL App..."
-python app.py

test_setup.py

@@ -1,54 +0,0 @@
-import sys
-import os
-from pathlib import Path
-import wfdb
-
-# Add camel_inference to path just like app.py does
-repo_dir = Path("./camel_inference").resolve()
-if str(repo_dir) not in sys.path:
-    sys.path.append(str(repo_dir))
-
-print("Testing imports...")
-try:
-    import gradio
-    import numpy
-    import matplotlib
-    import torch
-    print("Standard deps: OK")
-except ImportError as e:
-    print(f"Standard deps failed: {e}")
-    sys.exit(1)
-
-print("Testing CAMEL import...")
-try:
-    from camel.camel_model import CAMEL
-    print("CAMEL import: OK")
-except ImportError as e:
-    print(f"CAMEL import failed: {e}")
-    # Try looking in src explicitly if package install failed to link correctly
-    sys.path.append(str(repo_dir / "src"))
-    try:
-        from camel.camel_model import CAMEL
-        print("CAMEL import (via src fallback): OK")
-    except ImportError as e2:
-        print(f"CAMEL import totally failed: {e2}")
-        sys.exit(1)
-
-print("Testing WFDB Data Loading...")
-demo_dir = repo_dir / "demo"
-hea_files = list(demo_dir.glob("*.hea"))
-if not hea_files:
-    print("No .hea files found in demo dir!")
-    sys.exit(1)
-
-example_path = str(hea_files[0].with_suffix(""))
-try:
-    record = wfdb.rdrecord(example_path)
-    print(f"Loaded {record.record_name}: Fs={record.fs}, Shape={record.p_signal.shape}")
-except Exception as e:
-    print(f"WFDB load failed: {e}")
-    sys.exit(1)
-
-print("\n------------------------------")
-print("VERIFICATION SUCCESSFUL")
-print("------------------------------")