Fertility Medical Reasoner

This repository provides a lightweight fine-tuned medical reasoning model specialized for fertility and reproductive-medicine cases.
All demos are available in notebook/demo.ipynb.

1. Requirements

Install dependencies:

pip install torch transformers datasets peft trl unsloth huggingface_hub tqdm jupyterlab

Login to Hugging Face if your base model requires authentication:

from huggingface_hub import notebook_login
notebook_login()

2. Data Preparation

train/00_prepare_mixture.py is responsible for:

• Cleaning and preprocessing raw datasets
• Combining multiple sources into a unified training mixture
• Formatting examples for chain-of-thought supervised fine-tuning

Running this script will generate:

data/processed/final_mix/

3. Model Fine-Tuning

train/qlora_sft.py performs QLoRA-based supervised fine-tuning.

Key details:

• Base model: SmolLM3-3B-Base
• LoRA adapters for efficient training
• Long-context support (up to 32k tokens)
• Final model saved under:

models/sft_smollm3/

4. Evaluation

Two evaluation scripts are provided:

evaluation/run_all_benchmarks.py

Evaluates the fine-tuned model on medical-reasoning benchmarks:

• MedQA
• MedMCQA
• PubMedQA
• MMLU Clinical subsets

evaluation/run_batch_eval_base_model.py

Evaluates the base model on the same datasets for comparison.

Both scripts report:

• Accuracy per dataset
• Overall aggregated score
• Differences between yes/no reasoning and multiple-choice grounding

5. Inference and Use Cases

Use:

inference/stream_medical_cot.py

This script supports:

• Loading the fine-tuned model
• Passing new patient cases (JSON or text)
• Viewing chain-of-thought reasoning traces
• Getting recommended clinical actions

Useful for testing custom fertility cases or PDF-extracted patient summaries.

6. Demo

A complete end-to-end demo is available in:

notebook/demo.ipynb

It shows:

• Loading the fine-tuned model
• Running benchmark evaluations
• Testing inference on sample fertility cases
• Visualizing chain-of-thought reasoning