birgermoell/medqa-reasoning-traces

MedQA with LLM Reasoning Traces (kimi-k3)

A derived dataset pairing every MedQA USMLE question with a reasoning trace produced by a large language model, the model's predicted answer, and a correctness flag scored against the gold label.

• —Source questions: bigbio/med_qa, subset med_qa_en_4options_source (English, 4-option USMLE variant), splits train/validation/test
• —Records: 12,723 (one per question × model)
• —Overall accuracy: 94.4% (12,011/12,723)

Why this dataset?

MedQA measures clinical knowledge but ships without explanations of how an answer is reached. This dataset augments it with model-generated chain-of-thought, making it useful for:

• —Distillation / SFT data — (question, reasoning, answer) triples for training medical reasoning models
• —Reasoning analysis — studying how LLMs work through USMLE-style clinical cases
• —Evaluation — reproducible baseline runs with per-question correctness flags
• —Error analysis — filter to correct: false and inspect the flawed reasoning directly

Method

Source data

MedQA (Jin et al., 2021) is assembled from professional medical board exam materials. We used the med_qa_en_4options_source configuration of the bigbio/med_qa loader:

Each question carries its gold answer (answer_idx one of A–D) and all answer options.

Model and prompting

• —Model: kimi-k3, served through an OpenAI-compatible /v1/chat/completions endpoint
• —System prompt: You are an expert medical practitioner taking a board certification exam. Reason carefully through the clinical case, then give your final answer.
• —User prompt template:

Answer the following USMLE-style multiple-choice question.

{question}

(A) ...
(B) ...
(C) ...
(D) ...

First think through the problem carefully (step-by-step clinical reasoning).
Then end your response with exactly one final line in this format:

Answer: (X)

where X is the letter of the correct option.

• —Sampling: temperature 0.0; max_tokens 4,096 (records cut off by the cap were re-queried with max_tokens 8,192 and the longer response kept)
• —Concurrency: 8 parallel requests; each response flushed to disk the instant it completed

Reasoning trace extraction

The reasoning trace is captured from the reasoning field returned by the serving stack (vLLM prefix parsing for reasoning models). The extraction layer additionally supports reasoning_content (DeepSeek R1 style) and inline think-tag blocks inside content, for portability across providers.

Answer parsing and correctness

The final answer letter is parsed from the model's response text, in priority order:

1. 1.Explicit statement: The answer is (X) / Answer: X (case-insensitive, */parenthesis tolerant)
2. 2.A final bare line containing only a letter (E, **B**, …)
3. 3.Fallback: unique option-text match against the response body

The parsed letter is compared case-insensitively with the gold answer_idx. Records where no answer could be parsed are labeled correct: false with predicted_answer_idx: null.

Statistics

Data schema

Every record is one (question, model) pair:

Usage

from datasets import load_dataset

ds = load_dataset("birgermoell/medqa-reasoning-traces")
ex = ds["test"][0]
print(ex["gold_answer_idx"], ex["predicted_answer_idx"], ex["correct"])
print(ex["reasoning"][:500])

Filter for reasoning on incorrect answers:

wrong = ds["test"].filter(lambda x: not x["correct"])
for ex in wrong.select(range(3)):
    print(ex["question"][:120], "->", ex["gold_answer_idx"])

Known limitations and caveats

• —LLM reasoning is not human reasoning. Traces are model-generated and may contain plausible-but-wrong clinical claims, hallucinated facts, or post-hoc rationalizations — even when the final answer is correct. Do not treat this as medical content.
• —Correctness is string/letter-matched, not clinically adjudicated; gold labels come from MedQA itself and are not error-free.
• —Answer parsing follows a deterministic heuristic; ambiguous final answers default to correct: false rather than best-guess matching.
• —A small share of first-pass generations hit the 4,096-token cap and were re-queried at 8,192 tokens (the finish_reason of the kept record is stop unless even 8,192 was insufficient).
• —The reasoning trace reflects this specific serving stack's reasoning field for this model version; re-running with a newer model changes the traces.

Intended use

Research on medical QA, reasoning-trace collection, LLM evaluation and distillation. Not for clinical decision-making.

Licensing and attribution

• —Questions: MedQA (bigbio/med_qa) — assembled from USMLE practice materials; check the source dataset card before commercial use.
• —Reasoning traces: generated by kimi-k3; model output terms apply.

If you use this dataset, please also cite MedQA:

@inproceedings{jin2021medqa,
  title     = {What Disease does this Patient Have? A Large-scale Open Domain Question Answering Dataset from Medical Exams},
  author    = {Jin, Di and Pan, Eriol and Oufattole, Nassim and Weng, Wei-Hung and Fang, Hanyi and Szolovits, Peter},
  booktitle = {Findings of the Association for Computational Linguistics: EMNLP 2021},
  year      = {2021},
  doi       = {10.18653/v1/2021.findings-emnlp.437}
}

Dataset creation pipeline

The full pipeline is open source: questions are downloaded and normalized (prepare_dataset.py), queried with durable per-record writes and automatic retries (run_inference.py), token-capped records re-queried (--retry-length), joined into the derived dataset (build_derived.py), and served over HTTP.