论文解读
CLARITY主要实现的是模拟各种治疗方法和情境因素下,当前状况的患者的疾病发展变化预测。
CLARITY的核心设计原则之一,就是明确地将时间和临床环境作为分析的基础。连续的时间间隔被编码为嵌入形式,以便区分短期反应和长期演变。同时,还会整合患者的多种属性信息,包括基因组学、人口统计学以及治疗历史等,从而实现个性化的病情建模。这种设计使得模型能够应对现实临床数据中常见的不规则随访间隔和多样化的患者情况。
不同于MeWM使用扩散模型来合成治疗后的肿瘤图像,CLARITY更注重基于患者特定的临床信息来定制预测和治疗方案,并做患者的生存评估。
CLARITY 系统由三个主要组件构成:一个参数高效的视觉主干网络(MRI 编码器)、一个基于多模态大语言模型实现的疗法策略代理,以及一个行为模块(疾病演化模型)。此处需要注意,MRI编码器是预训练的,且权重冻结。因此,本论文复现时需要重复训练的仅为疾病演化模型。
关于数据集方面,本文选用了三个MRI数据集,其中两个为脑肿瘤垂类数据集:
名称 | 类型 | 作用 |
|---|---|---|
MU-Glioma-Post | 脑肿瘤数据集 | 初始训练。75%的数据用于训练,15%用于验证,15%用于测试。 |
UCSF-ALPTDG | 脑肿瘤数据集 | 零样本外部验证。 |
ISPY-2 | 乳腺癌数据集 | 模型评估。独立训练。 |
代码克隆
首先克隆下CLARITY的实现代码:
BASH
git clone https://ghproxy.net/https://github.com/DingTianxingjian/CLARITY.git数据
医学影像
病人的医学影像来源于TCIA的MU-Glioma-Post数据集,前往此处下载(需要使用IBM-Aspera-Connect plugin)
下载后移动到./CLARITY/datasets文件夹下,形成如下文件夹结构:
PLAINTEXT
root@user-System-Product-Name:/home/user/lsj/CLARITY# tree
.
├── datasets
│ └── MU-Glioma-Post
│ ├── PatientID_0003
│ │ ├── Timepoint_1
│ │ │ ├── PatientID_0003_Timepoint_1_brain_t1c.nii.gz
│ │ │ ├── PatientID_0003_Timepoint_1_brain_t1n.nii.gz
│ │ │ ├── PatientID_0003_Timepoint_1_brain_t2f.nii.gz
│ │ │ ├── PatientID_0003_Timepoint_1_brain_t2w.nii.gz
│ │ │ └── PatientID_0003_Timepoint_1_tumorMask.nii.gz
│ │ ├── Timepoint_2
│ │ │ ├── PatientID_0003_Timepoint_2_brain_t1c.nii.gz
│ │ │ ├── PatientID_0003_Timepoint_2_brain_t1n.nii.gz
│ │ │ ├── PatientID_0003_Timepoint_2_brain_t2f.nii.gz
│ │ │ ├── PatientID_0003_Timepoint_2_brain_t2w.nii.gz
│ │ │ └── PatientID_0003_Timepoint_2_tumorMask.nii.gz注意:Firefox浏览器对下载插件的兼容性较差,可能导致下载失败
预训练模型
从Dropbox(显然要科学上网)下载BrainIAC的预训练权重,并放置在BrainIAC-main/src/checkpoints/BrainIAC.ckpt ,注意文件如下图:
配置
该代码基于Python3.10环境运行,首先新建环境:
PLAINTEXT
conda create -n clpy310 python=3.10 -c conda-forge然后激活环境:
PLAINTEXT
conda activate clpy310然后安装依赖:
PLAINTEXT
pip install -r ./requirements.txt有些重复的造成错误,自行删减一下。
依旧有经典几个包重新配一下:
PLAINTEXT
conda install conda-forge::pyairports
pip install git+https://ghproxy.net/https://github.com/Radiomics/pyradiomics.git@v3.1.0
python -m pip install "setuptools<82.0.0" --force-reinstall有几个版本错了:
PLAINTEXT
pip install -U bitsandbytes>=0.46.1
pip install 'accelerate>=1.1.0'然后根据nvidia-smi选择安装torch:
PLAINTEXT
pip install torch torchvision torchaudio -f https://mirrors.aliyun.com/pytorch-wheels/cu132/训练
首先因为国内网络环境的原因,先设定HF的镜像地址:
SHELLSCRIPT
export HF_ENDPOINT=https://hf-mirror.com接着使用CLARITY/Predictor/dataset/extract_clinicial.py 转换数据集里的Excel文件为features_csv(Predictor/dataset/MU_Glioma_Post/features_output.csv)
如果出现:
Cannot access gated repo for url https://hf-mirror.com/google/medgemma-4b-it/resolve/main/config.json.
Access to model google/medgemma-4b-it is restricted and you are not in the authorized list. Visit https://huggingface.co/google/medgemma-4b-it to ask for access.
需要去链接里点一下授权,然后再使用hf auth login 登录。
然后开始训练即可:
PLAINTEXT
python Predictor/train.py \
--exp_name exp012_cf_diversity \
--features_csv Predictor/dataset/MU_Glioma_Post/features_output.csv \
--timeline_json Predictor/dataset/MU_Glioma_Post/clinical_latest.json \
--mri_data_dir datasets/MU-Glioma-Post \
--brainiac_ckpt BrainIAC-main/src/checkpoints/BrainIAC.ckpt \
--brainiac_tokens 8 \
--brainiac_lora_r 8 \
--latent_dim 768 \
--num_modalities 32 \
--take_dims 767 \
--text_encoder_name google/medgemma-4b-it \
--num_epochs 100 \
--batch_size 4 \
--val_batch_size 4 \
--num_workers 4 \
--lr 2e-4 \
--text_lr 1e-4 \
--text_proj_lr 5e-4 \
--cf_weight 1.0 \
--cf_cos_margin 0.9 \
--lambda_l1 0.5 \
--lambda_cox 1.0 \
--lambda_bce 1.0 \
--dropout 0.3 \
--grad_clip_norm 2.0 \
--survival_wd 0.01 \
--warmup_epochs 10 \
--seed 42等待训练结束即可,这个代码跑得比MeWM好很多。
验证
需要修改评估脚本以支持Deepseek,然后进行评估:
SHELLSCRIPT
# 1. 获取 DeepSeek API Key: https://platform.deepseek.com/
export OPENAI_API_KEY="sk-****"
# 2. 指定模型和端点
export ISE_LLM_MODEL="deepseek-v4-flash"
export ISE_LLM_PROVIDER="openai"
export ISE_LLM_BASE_URL="https://api.deepseek.com/"
# 3. 运行评估
python ise_eval_v2.py得到评估结果:
PLAINTEXT
======================================================================
ISE vs LLM F1 (N=24 patients, w_surv=0.5)
======================================================================
Method F1 mean F1 std F1 median
-------------------------------------------------------
Random 0.3069 0.3944 0.0000
LLM primary 0.4194 0.4191 0.4500
ISE (WM+LLM) 0.4292 0.3892 0.5000
Oracle (UB) 0.5736 0.4287 0.6667
Exact match ISE: 0.167 | LLM: 0.250
ISE > LLM: 3 | ISE = LLM: 16 | ISE < LLM: 5
ISE Δ F1 vs LLM: +0.0097
Per-phase F1 (ISE / LLM):
gbm_chemonaive n= 9 ISE=0.330 LLM=0.341 Δ=-0.011
other n= 2 ISE=1.000 LLM=0.500 Δ=+0.500
post_crt n=13 ISE=0.410 LLM=0.462 Δ=-0.051
GT distribution: [('TMZ', 14), ('BEV', 13), ('TTF', 5), ('RT', 4)]
Results saved to ise_eval_v2_results.json总体来看比大模型好一点,主要是样本量太小了,存在过拟合。
所有代码修改
主要是适配双卡训练和适配Deepseek接入的修改
PLAINTEXT
diff --git a/Predictor/models/full_model.py b/Predictor/models/full_model.py
index 07b6695..41bdc22 100644
--- a/Predictor/models/full_model.py
+++ b/Predictor/models/full_model.py
@@ -443,9 +443,11 @@ class SharedTextEncoder(nn.Module):
load_kwargs = dict(
dtype=torch.bfloat16,
- device_map="auto",
low_cpu_mem_usage=True,
)
+ # 单卡加载:不传 device_map。
+ # bitsandbytes 4-bit 量化会自动将模型放到 GPU:0,避免 accelerate hooks
+ # 在 PEFT 多卡场景下的跨设备冲突。
if bnb_config is not None:
load_kwargs["quantization_config"] = bnb_config
@@ -530,9 +532,7 @@ class SharedTextEncoder(nn.Module):
else:
emb = hidden_states[:, 0, :]
-
-
- emb = self.final_proj(emb.to(self.final_proj.weight.device))
+ emb = self.final_proj(emb)
emb = F.normalize(emb, p=2, dim=-1)
return emb
diff --git a/ise_eval_v2.py b/ise_eval_v2.py
index 1090fb4..2319c68 100644
:
diff --git a/Predictor/models/full_model.py b/Predictor/models/full_model.py
index 07b6695..41bdc22 100644
--- a/Predictor/models/full_model.py
+++ b/Predictor/models/full_model.py
@@ -443,9 +443,11 @@ class SharedTextEncoder(nn.Module):
load_kwargs = dict(
dtype=torch.bfloat16,
- device_map="auto",
low_cpu_mem_usage=True,
)
+ # 单卡加载:不传 device_map。
+ # bitsandbytes 4-bit 量化会自动将模型放到 GPU:0,避免 accelerate hooks
+ # 在 PEFT 多卡场景下的跨设备冲突。
if bnb_config is not None:
load_kwargs["quantization_config"] = bnb_config
@@ -530,9 +532,7 @@ class SharedTextEncoder(nn.Module):
else:
emb = hidden_states[:, 0, :]
-
-
- emb = self.final_proj(emb.to(self.final_proj.weight.device))
+ emb = self.final_proj(emb)
emb = F.normalize(emb, p=2, dim=-1)
return emb
diff --git a/ise_eval_v2.py b/ise_eval_v2.py
index 1090fb4..2319c68 100644
--- a/ise_eval_v2.py
+++ b/ise_eval_v2.py
@@ -37,22 +37,29 @@ DISCOUNT = 0.95
ROLLOUT_N = 5
DEVICE = torch.device("cuda")
-# Set ANTHROPIC_API_KEY or OPENAI_API_KEY to use real LLM candidates.
-# Without a key the rule engine fallback runs automatically.
-_API_KEY = os.environ.get("ANTHROPIC_API_KEY") or os.environ.get("OPENAI_API_KEY")
-_LLM_MODEL = os.environ.get("ISE_LLM_MODEL", "claude-opus-4-8")
-_LLM_PROVIDER = "anthropic" if os.environ.get("ANTHROPIC_API_KEY") else "openai"
+# LLM provider configuration — supports Anthropic, OpenAI, and OpenAI-compatible (DeepSeek, etc.)
+# DeepSeek example:
+# export OPENAI_API_KEY="sk-your-deepseek-key"
+# export ISE_LLM_MODEL="deepseek-chat"
+# export ISE_LLM_PROVIDER="openai"
+# export ISE_LLM_BASE_URL="https://api.deepseek.com"
+_API_KEY = os.environ.get("ANTHROPIC_API_KEY") or os.environ.get("OPENAI_API_KEY")
+_LLM_MODEL = os.environ.get("ISE_LLM_MODEL", "claude-opus-4-8")
+_LLM_PROVIDER = os.environ.get("ISE_LLM_PROVIDER",
+ "anthropic" if os.environ.get("ANTHROPIC_API_KEY") else "openai")
+_LLM_BASE_URL = os.environ.get("ISE_LLM_BASE_URL", None)
try:
from ise_llm import ISECandidateGenerator
_llm_gen = ISECandidateGenerator(
- api_key=_API_KEY, model=_LLM_MODEL, provider=_LLM_PROVIDER
+ api_key=_API_KEY, model=_LLM_MODEL, provider=_LLM_PROVIDER,
+ base_url=_LLM_BASE_URL,
) if _API_KEY else None
except Exception as _e:
print(f"[ise_llm] not available ({_e}), using rule engine fallback.")
_llm_gen = None
-print(f"LLM backend: {'API (' + _LLM_MODEL + ')' if _llm_gen else 'rule engine (no API key)'}")
+print(f"LLM backend: {'API (' + _LLM_MODEL + ' @ ' + _LLM_PROVIDER + ')' if _llm_gen else 'rule engine (no API key)'}")
# ═══════════════════════════════════════════════════════════════════
# Treatment fingerprint (for F1)
@@ -286,15 +293,26 @@ print("=" * 70)
print("Loading exp012 model ...")
args = json.load(open(ARGS_JSON))
+# Build vision backbone (same as train.py:build_model)
+vision_backbone = None
+brainiac_ckpt_path = args.get("brainiac_ckpt")
+if brainiac_ckpt_path:
+ from models.brainiac_adapter import BrainIACAdapter
+ from models.vision_backbone import MultiModalVisionBackbone
+ _adapter = BrainIACAdapter(
+ checkpoint_path=brainiac_ckpt_path,
+ tokens_per_modality=args.get("brainiac_tokens", 8),
+ lora_r=args.get("brainiac_lora_r", 8),
+ )
+ vision_backbone = MultiModalVisionBackbone(_adapter, num_modalities=4)
+
model = TimeAwareGliomaSurvivalPredictor(
text_encoder_name=args["text_encoder_name"],
latent_dim=args["latent_dim"],
num_modalities=args["num_modalities"],
predictor_hidden_dim=512,
survival_hidden_dim=128,
- brainiac_ckpt=args.get("brainiac_ckpt"),
- brainiac_tokens_per_modality=args.get("brainiac_tokens", 8),
- brainiac_lora_r=args.get("brainiac_lora_r", 8),
+ vision_backbone=vision_backbone,
).to(DEVICE)
ckpt = torch.load(CKPT, map_location=DEVICE, weights_only=False)
diff --git a/ise_llm.py b/ise_llm.py
index b56bc8e..9e5c73d 100644
--- a/ise_llm.py
+++ b/ise_llm.py
@@ -76,6 +76,8 @@ scored poorly, or escalate if low-intensity plans scored poorly relative to alte
```
Omit empty arrays. "actions": {} is valid for observation-only candidates.
Valid action keys: radiation, chemotherapy, immunotherapy, additional_1, additional_2, other_therapy.
+
+CRITICAL: Output ONLY valid JSON inside a ```json code block. No greetings, no explanations, no markdown outside the code block. Start directly with ```json and end with ```.
"""
# ── User prompt builder ───────────────────────────────────────────────────────
@@ -128,16 +130,38 @@ def _build_user_prompt(patient_context: dict, pre_actions: dict,
# ── JSON parser ───────────────────────────────────────────────────────────────
def _parse_candidates(raw_text: str) -> list:
- """Extract the JSON block from LLM output and parse candidates."""
- # Try to find ```json ... ``` block
- m = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", raw_text, re.DOTALL)
+ """Extract the JSON block from LLM output and parse candidates.
+ Handles: markdown code blocks, surrounding text, DeepSeek thinking output."""
+ # 1. Try ```json ... ``` code block (greedy, across newlines)
+ m = re.search(r"```(?:json)?\s*(\{.*\})\s*```", raw_text, re.DOTALL)
if m:
raw_text = m.group(1)
- # Fall back: first { ... } block
- m = re.search(r"\{.*\}", raw_text, re.DOTALL)
- if not m:
+
+ # 2. Find outermost { ... } by brace counting (robust to surrounding text)
+ start = raw_text.find("{")
+ if start == -1:
raise ValueError(f"No JSON found in LLM output:\n{raw_text[:500]}")
- data = json.loads(m.group(0))
+ depth = 0
+ end = -1
+ for i in range(start, len(raw_text)):
+ if raw_text[i] == "{":
+ depth += 1
+ elif raw_text[i] == "}":
+ depth -= 1
+ if depth == 0:
+ end = i + 1
+ break
+ if end == -1:
+ raise ValueError(f"Unbalanced braces in LLM output:\n{raw_text[:500]}")
+
+ json_str = raw_text[start:end]
+ try:
+ data = json.loads(json_str)
+ except json.JSONDecodeError:
+ # Last resort: try to fix trailing commas, single quotes etc.
+ json_str = re.sub(r",\s*}", "}", json_str)
+ json_str = re.sub(r",\s*]", "]", json_str)
+ data = json.loads(json_str)
return data.get("candidates", [])
@@ -167,10 +191,12 @@ class ISECandidateGenerator:
def __init__(self, api_key: Optional[str] = None,
model: str = "claude-opus-4-8",
provider: str = "anthropic",
+ base_url: Optional[str] = None,
max_tokens: int = 2048,
temperature: float = 0.3):
self.model = model
self.provider = provider
+ self.base_url = base_url
self.max_tokens = max_tokens
self.temperature = temperature
@@ -183,7 +209,10 @@ class ISECandidateGenerator:
self._client = anthropic.Anthropic(api_key=key)
else:
import openai
- self._client = openai.OpenAI(api_key=key)
+ kwargs = dict(api_key=key)
+ if base_url:
+ kwargs["base_url"] = base_url
+ self._client = openai.OpenAI(**kwargs)
# ── Low-level call ──────────────────────────────────────────────────
@@ -197,7 +226,7 @@ class ISECandidateGenerator:
messages=[{"role": "user", "content": user_message}],
)
return resp.content[0].text
- else: # openai-compatible
+ else: # openai-compatible (DeepSeek, OpenAI, etc.)
resp = self._client.chat.completions.create(
model=self.model,
max_tokens=self.max_tokens,
@@ -206,6 +235,7 @@ class ISECandidateGenerator:
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": user_message},
],
+ response_format={"type": "json_object"},
)
return resp.choices[0].message.content