ClinSeek-35B-A3B-DeepmedSFT-Think

ClinSeek-35B-A3B-DeepmedSFT-Think

Full-parameter SFT checkpoints of Qwen/Qwen3.5-35B-A3B on agentic EHR trajectories (Chtholly17/Deepmed_SFT), trained with the ClinSeekAgent SFT recipe (https://github.com/UCSC-VLAA/ClinSeekAgent, vendored verl + Megatron/mbridge).

The distinguishing feature of this run: the teacher's forced ehr.think tool calls were folded into native <think>…</think> reasoning before every action turn, so the student learns to reason in the Qwen3.5 thinking channel instead of emitting a think tool call.

Checkpoints in this repo

Three training runs share this repo. They use the same recipe and preprocessing; they differ in which part of Chtholly17/Deepmed_SFT they train on (runs A/B: AgentEHR-Bench prompts; run C: EHRBench prompts, a different task mix — see Training data).

Validation sets differ between the runs (run A: 53 random rows; run B: 108 rows from 56 held-out questions; run C: 112 rows from 112 held-out questions), so the val losses are only roughly comparable across runs. Within each run the third epoch lowers train loss to ~0.12-0.13 while val loss rises (A: 0.316 -> 0.331, B: 0.274 -> 0.291, C: 0.305 -> 0.317), so the epoch-2 checkpoints are the ones to use. The epoch-1 checkpoints of all runs and the epoch-3 checkpoints of runs B and C were not uploaded.

Each subfolder is a complete HF checkpoint (Qwen3_5MoeForConditionalGeneration, bf16, 14 safetensors shards, tokenizer, chat_template.jinja, generation_config.json).

Training data

• Source: Chtholly17/Deepmed_SFT — think-first agentic EHR trajectories generated by gpt-5.6-sol (reasoning effort high), EHR tools only (16 ehr.* MCP tools, no browser), temporal policy strict-ehr-runtime-v8. The dataset covers two prompt sets:
  • AgentEHR3k (AgentEHR3k_..._run1.jsonl, 2,947 successful trajectories, and ..._run2.jsonl, 2,972): the 3,000 prompts of the AgentEHR-Bench train split, six task families (diagnoses_ccs, procedures_ccs, labevents, microbiologyevents, prescriptions, transfers), answered with a list of official candidate names. Two independent sampling passes over the same prompts.
    • Run A (epoch2_step164, epoch3_step246) used run1 only.
    • Run B (run12_epoch2_step330) used run1 + run2 (5,919 trajectories; 2,920 prompts carry two different trajectories, 79 carry one).
  • EHRBench6k (EHRBench6k_gpt56sol_thinkfirst_strictv8_run1.jsonl, 5,984 successful trajectories, one per prompt): 45 EHRBench task types in two families. 18 risk-prediction tasks answered yes/no (ED hospitalization, critical outcomes, ICU transfer within 12 h, ED re-attendance, inpatient and ICU mortality at 1/2/3/7/14 days, ICU stay > 7/14 days, ICU readmission, length of stay > 3/7 days, readmission within 30/60 days; 2,999 trajectories), and 27 decision-making tasks that pick the next clinical event or code from a <candidate_answers> list given in the prompt (next diagnosis as ICD/CCS, procedures, prescriptions/ATC, lab / microbiology / chart / input / output / datetime events, medication reconciliation, transfers, services, radiology, admissions, POE, Pyxis, OMR; 2,985 trajectories). Trajectories average 13.9 assistant turns (max 26).
    • Run C (ehrbench6k_epoch2_step344) used this file only.
• Preprocessing (prepare_clinseek_data.py --success_only --merge_think --strip_think_protocol):
  1. Every assistant turn that only calls ehr.think is removed together with its "Thinking Finish" tool response; the response argument text becomes the reasoning_content of the next assistant turn, which the Qwen3.5 chat template renders as <think>…</think>. Consecutive think turns are concatenated (runs A/B: 26,428 think turns folded; run C: 41,707; 0 ehr.think calls remain).
  2. The REASONING TRACE PROTOCOL paragraph (which instructed the teacher to call ehr.think first in every round) is stripped from the system prompt. The - ehr.think: Record your reasoning process line in the tool list is kept, matching the deployment prompt.
  3. Tool-call arguments are kept as structured dicts so the template renders native <tool_call> / <tool_response> blocks; the provider's own reasoning_content summaries are discarded.
  4. Samples longer than 52,000 tokens (Qwen3.5 template) are dropped (run1 2,681 / 2,947; run1+run2 5,394 / 5,919; EHRBench6k 5,637 / 5,984 — 91-94% kept).
  5. Train/val split with seed 42. Run A: random rows -> 2,628 train / 53 val. Run B: split by question id (both trajectories of a prompt land in the same split, zero qid overlap) -> 5,286 train / 108 val (56 held-out questions). Run C: split by question id -> 5,525 train / 112 val (112 held-out questions).
• Loss is computed on assistant tokens only (thinking block + tool calls + final answer); system, user and tool-response tokens are masked.

Prompt format

System prompt (verbatim, identical for every training sample)

You are a research assistant with access to clinical EHR tools.

**EHR Tools** (for clinical data analysis):
- ehr.load_ehr: Load patient EHR database (must be called first for clinical tasks)
- ehr.get_table_names: List available patient data tables
- ehr.get_column_names: Get table column information
- ehr.get_records_by_time: Query patient records within time range
- ehr.run_sql_query: Execute SQL queries on patient database
- ehr.get_candidates_by_semantic_similarity: Search medical terminology/diagnosis codes
- ehr.get_candidates_by_keyword: Search diagnosis codes by keyword
- ehr.think: Record your reasoning process
- ehr.finish: Submit your final answer

**Important:** No web browsing is available in this setting. Base every
conclusion on the patient data returned by the EHR tools and on your own
clinical knowledge; do not attempt to search the web.

sources=ehr

This is the teacher's system prompt with the REASONING TRACE PROTOCOL paragraph removed. Only the seven tools above plus ehr.think/ehr.finish are listed, but the trajectories use all 16 ehr.* MCP tools (e.g. ehr.get_latest_records, ehr.get_candidates_by_fuzzy_matching, ehr.get_records_by_keyword, ehr.get_event_counts_by_time, ehr.get_table_description, ehr.get_unique_values, ehr.get_records_by_value). No tools= JSON schema is passed to the chat template; the model learned the tool names and argument schemas from the trajectories. The model never calls ehr.think (0 occurrences in the training targets) — its reasoning goes into the <think> block instead.

User turn (AgentEHR-Bench task prompt)

<task_instruction>
Your current task is to act as a <role, e.g. clinical microbiologist / diagnostician>.

Your objective is to determine all <targets> for the patient by analyzing their complete
medical history ...

You must find the most likely official <...> candidates using the **`<task>_candidates`**
reference data or semantic matching tools.
...
Submit your final answer through `ehr.finish` as a **list** containing **multiple plausible
<items>**. Each item in the list must be a string representing an official <...> name, and
**must not contain any codes or other additional information**.
</task_instruction>

<patient_info>
Current Time: 2155-08-29 23:59:00
Patient Subject ID: 19573105
</patient_info>

EHRBench6k prompts (run C) have the same three-part shape, with a shorter instruction and a candidate list:

<task_instruction>
Please give the next ED Diagnoses on International Classification of Diseases suggestion for the patiens.

Analyze the patient's EHR records up to the current time to make your prediction.
Do not assume any unseen future events beyond the current time.
When you need medical knowledge or clinical information, use the `browser.search` tool to find authoritative medical information from reliable sources.
Submit your final answer through `ehr.finish`. Your answer must be chosen only from the candidate list below.
</task_instruction>

<patient_info>
Current Time: 2132-05-28 15:15:01
Patient Subject ID: 18185480
</patient_info>

<candidate_answers>
["Oth transient cerebral ischemic attacks and related synd", "Other epilepsy, not intractable, without status epilepticus", ...]
</candidate_answers>

Risk-prediction prompts ask a yes/no question instead (e.g. whether the patient will be admitted from the ED, or die in the ICU within 7 days) and the candidate list is ["yes", "no"].

Note: the benchmark task text still mentions browser.search; the teacher ignored it (no browser tool was available, as the system prompt states) and so does the student.

Output format

Every assistant turn produced by the model has the shape below (Qwen3.5 chat template, content is always empty in the training targets — 0 of 2,659 assistant turns contain free text outside the think block and tool calls):

<|im_start|>assistant
<think>
The available EHR lacks dedicated diagnosis, procedure, note, radiology, and microbiology
event tables, but includes admissions/ED data, labs, medication administration, orders,
transfers, triage, and vital signs. Therefore infection syndromes must be reconstructed from
... I will inspect schemas for admissions, ED stays, triage, vitals, labs, orders,
medication administration, and transfers in parallel; ...
</think>

<tool_call>
<function=ehr.get_column_names>
<parameter=subject_id>
19573105
</parameter>
<parameter=table_name>
admissions
</parameter>
</function>
</tool_call>
<tool_call>
<function=ehr.get_column_names>
<parameter=subject_id>
19573105
</parameter>
<parameter=table_name>
edstays
</parameter>
</function>
</tool_call><|im_end|>

• The <think> block is the folded teacher reasoning: what the previous results showed, which hypotheses they raise or rule out, and what to query next. It is present in every assistant turn (including the first one and the final answer turn).
• One or several <tool_call> blocks follow in the same turn (median 1, up to 8 parallel calls), in Qwen3.5's XML-style function-call syntax (<function=name> / <parameter=key>value</parameter>), which vLLM's qwen3_xml tool parser converts to OpenAI tool_calls.
• Tool results are fed back as user turns wrapped in <tool_response> ... </tool_response> (one block per call, in call order), exactly as the Qwen3.5 template renders role: tool messages.
• The trajectory ends with a single call to ehr.finish whose only argument is response, a JSON list of official candidate names (runs A/B), or, for run C, ["yes"] / ["no"] for risk-prediction tasks and a list of names taken from <candidate_answers> for decision-making tasks:

<think>
... Given the absence of any clinical infection syndrome, invasive cultures, C. difficile
testing, MRSA screening, and serologies would be unnecessary. I will therefore submit only
exact official names for plausible noninvasive respiratory viral screening ...
</think>

<tool_call>
<function=ehr.finish>
<parameter=response>
["Respiratory Viral Antigen Screen", "Respiratory Virus Identification", "DIRECT INFLUENZA A ANTIGEN TEST", "DIRECT INFLUENZA B ANTIGEN TEST", "DIRECT RSV ANTIGEN TEST"]
</parameter>
</function>
</tool_call>

Because the model reasons in the thinking channel, serve it with a reasoning parser (vLLM: --reasoning-parser qwen3) so the <think> text is returned as reasoning_content and not mixed into content; use --tool-call-parser qwen3_xml for the tool calls.

Training configuration (shared by all checkpoints)

Differences from the ClinSeekAgent paper recipe: 2 GPUs instead of 8 (TP2/EP2 instead of TP2/EP8), a different and smaller dataset (Deepmed_SFT think-first trajectories, EHR tools only, instead of the Claude Opus trajectories with browser tools), and the think-folding preprocessing above. Learning rate, batch size, schedule and sequence length follow the paper.

Loss curves

Run B (run1 + run2, 5,286 train samples, 165 steps/epoch)

Run C (EHRBench6k, 5,525 train samples, 172 steps/epoch)

Run C was resumed twice (after step 172 and after step 344, see Caveats); steps 173-344 and 345-516 come from the resumed processes.

Run A (run1 only, 2,628 train samples, 82 steps/epoch)

Train loss (mean over 41-step blocks) and validation loss:

Usage

from transformers import AutoTokenizer, AutoModelForCausalLM
repo = "Chtholly17/ClinSeek-35B-A3B-DeepmedSFT-Think"
tok = AutoTokenizer.from_pretrained(repo, subfolder="run12_epoch2_step330")
model = AutoModelForCausalLM.from_pretrained(repo, subfolder="run12_epoch2_step330", dtype="bfloat16", device_map="auto")

For vLLM, download the subfolder and serve it with the Qwen3 reasoning parser so the <think> block is separated from the tool-call content:

hf download Chtholly17/ClinSeek-35B-A3B-DeepmedSFT-Think --include "run12_epoch2_step330/*" --local-dir ./ckpt
vllm serve ./ckpt/run12_epoch2_step330 --reasoning-parser qwen3 --enable-auto-tool-choice --tool-call-parser qwen3_xml

See Prompt format and Output format above for the exact system prompt and the turn structure. Use the model with the ClinSeekAgent EHR MCP tools and the system prompt shown above (or the standard AgentEHR one); do not add the teacher's "call ehr.think first" instruction.

Caveats and license

• Trained on trajectories over MIMIC-IV-derived patient records (PhysioNet credentialed data). The weights are released for research use; users must comply with the PhysioNet data use agreement and must not attempt to recover patient-level information.
• Small datasets (2.6k / 5.3k / 5.5k trajectories): the third epoch overfits by validation loss in all runs; prefer run12_epoch2_step330 (or epoch2_step164 for the run1-only variant) for AgentEHR-Bench-style tasks and ehrbench6k_epoch2_step344 for EHRBench-style tasks unless downstream evaluation says otherwise.
• Run C was not a single uninterrupted process. The original process stalled at step 187 (page-cache thrashing on a shared node, no numerical problem); training was restarted from the epoch-1 checkpoint (step 172), and the restarted process stopped at step 344 because the restored dataloader state ended an epoch early, so the last epoch was started again from the epoch-2 checkpoint (step 344). At both resumes the model weights and the data position were restored and the cosine LR schedule was fast-forwarded to the correct step (with a 10-step linear re-warmup of the LR), but the Adam first/second moments were re-initialised because the checkpoints only hold model weights, and epoch 3 used a different shuffle seed than an uninterrupted run would have. Train/val losses after each resume line up with the values before it (e.g. val 0.315 at step 150 -> 0.310 at step 175), so we do not expect a material effect, but the run is not bit-for-bit equivalent to runs A/B.
• Not evaluated on AgentEHR-Bench yet; validation loss is the only quality signal reported here.
• Not a medical device. Outputs must not be used for clinical decision-making.