LLM fine-tuning on sensitive data

Load the clinical notes

The dataset has 1,028 training notes, 672 validation, 1,681 test. Each record is a free-text encounter note paired with the structured fields we'll teach the model to extract.

from datasets import load_dataset

ds = load_dataset("TonicAI/synthetic_clinical_notes")
example = ds["train"][0]
DatasetDict({
    train:      1028 rows  ['note', 'encounter_data']
    validation:  672 rows
    test:       1681 rows
})

De-identify with Tonic Textual

Textual detects PHI spans — names, dates, addresses, phone numbers, IDs — and replaces them. Two configuration choices matter for fine-tuning:

• Synthesis, not redaction. Replacing “Harvey D'Amore” with a [NAME_GIVEN] token would leave unnatural text and destroy the extraction labels. Synthesis swaps in a realistic substitute, so the notes still read like notes and the task is unchanged.
• Consistency between note and labels. Each {note, encounter_data} record goes through redact_json as a single unit with a fixed per-record seed, so the same synthetic name lands in the note text and in the ground-truth JSON. The input→output mapping the model has to learn survives de-identification intact.

Everything that isn't PHI — symptoms, procedures, medications, clinical reasoning — passes through untouched.

These are the entity types synthesized in this playbook (Textual detects many more — see the full entity library):

PHI_LABELS = [
    "NAME_GIVEN", "NAME_FAMILY", "LOCATION_ADDRESS", "LOCATION_STATE",
    "LOCATION_CITY", "LOCATION_ZIP", "LOCATION_COUNTRY", "PHONE_NUMBER",
    "EMAIL_ADDRESS", "CREDIT_CARD", "CC_EXP", "CVV", "MONEY", "ORGANIZATION",
    "DOB", "DATE_TIME", "URL", "NUMERIC_PII", "HEALTHCARE_ID",
]

DEID_CONFIG = dict(
    generator_config={label: "Synthesis" for label in PHI_LABELS},
    # keep synthetic given names gender-consistent — gender is an extraction target
    generator_metadata={"NAME_GIVEN": NameGeneratorMetadata(preserve_gender=True)},
    generator_default="Off",
)

response = textual.redact_json(example, random_seed=42, **DEID_CONFIG)
synthetic_example = json.loads(response.redacted_text)

One call, one record — here's what changes and what doesn't:

Now the full training split, with a distinct seed per record — synthetic identities vary across the dataset but stay consistent within each record. Only the training split gets de-identified, because it's the thing that leaves the compliance boundary; evaluation happens later against the real validation notes, which stay in your environment.

def synthesize_record(args):
    i, record = args
    response = textual.redact_json(record, random_seed=i, **DEID_CONFIG)
    synthetic = json.loads(response.redacted_text)
    # redact_json drops empty fields ([] / null); restore them so the
    # training labels keep the full schema — empty values carry no PHI
    for key, value in record["encounter_data"].items():
        if key not in synthetic["encounter_data"]:
            assert not value, f"non-empty field {key!r} went missing"
            synthetic["encounter_data"][key] = value
    return synthetic

with cf.ThreadPoolExecutor(max_workers=6) as ex:
    synthetic_train = list(tqdm(
        ex.map(synthesize_record, enumerate(ds["train"])),
        total=len(ds["train"]),
    ))

A quick audit before we call this safe: did any original identity survive into its synthetic counterpart?

original_names = [r["encounter_data"]["name"] for r in ds["train"]]
synthetic_names = [r["encounter_data"]["name"] for r in synthetic_train]

survivors = [
    orig for orig, synth in zip(original_names, synthetic_train)
    if orig in synth["note"] or orig == synth["encounter_data"]["name"]
]

print("records where the original name survived:", len(survivors))
print("full-name overlap between real and synthetic sets:",
      len(set(original_names) & set(synthetic_names)))
      
records where the original name survived: 0
full-name overlap between real and synthetic sets: 0

Define the extraction task

The target is one typed record per encounter. race and gender are constrained to HL7 value sets (US Core Race, Gender Identity); the remaining fields are free-form but formatted per rules in the system prompt. This isn't a full FHIR resource — it's the structured, vocabulary-controlled layer you'd map into one.

The Pydantic model does double duty: it documents the schema inside the system prompt, and at evaluation time it validates model outputs.

class GenderIdentityEnum(str, Enum):
    """HL7 THO 'Gender Identity' value set."""
    FEMALE = "female"
    MALE = "male"
    NON_BINARY = "non-binary"
    ASKED_DECLINED = "asked-declined"
    UNKNOWN = "unknown"

class FhirUsCoreRaceEnum(str, Enum):
    """HL7 'US Core Race' value set."""
    WHITE = "white"
    ASIAN = "asian"
    AFRICAN_AMERICAN = "black or african american"
    NATIVE_HAWAIIAN = "native hawaiian or other pacific islander"
    AMERICAN_INDIAN = "american indian or alaska native"
    UNKNOWN = "unknown"

class EncounterData(BaseModel):
    name_given: str              # patient's first given name
    name_family: str             # patient's family name (surname)
    age: int                     # age in years at encounter
    race: FhirUsCoreRaceEnum
    gender: GenderIdentityEnum
    reason: Optional[str]        # short diagnosis/complaint, title case
    procedures: List[str]        # distinct procedures, sentence case
    medications: List[str]       # generic drug names, [] if none
    encounter_type: str          # e.g. "encounter for check up"

Each record becomes a chat conversation — Fireworks fine-tunes on JSONL where every line is {"messages": [...]}, the same shape as a chat-completion request. The assistant target is projected through the schema, so the labels contain exactly the fields the model must produce. One file per training arm.

def to_messages(record):
    target = EncounterData.model_validate(record["encounter_data"]).model_dump(mode="json")
    return {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},   # full prompt in the notebook
        {"role": "user", "content": record["note"]},
        {"role": "assistant", "content": json.dumps(target)},
    ]}

write_jsonl(DATA_DIR / "train_real.jsonl", ds["train"])
write_jsonl(DATA_DIR / "train_synthetic.jsonl", synthetic_train)

Fine-tune on Fireworks

Two identical LoRA jobs on Llama 3.2 3B Instruct, differing only in their training data. Worth stating plainly: in a real deployment, only train_synthetic.jsonl would ever be uploaded. The real training file goes up here purely to build the baseline that measures what synthesis costs.

real_dataset  = ensure_dataset("clinical-notes-train-real", DATA_DIR / "train_real.jsonl")
synth_dataset = ensure_dataset("clinical-notes-train-synthetic", DATA_DIR / "train_synthetic.jsonl")

def ensure_sft_job(job_id, dataset_name):
    try:
        return fw.supervised_fine_tuning_jobs.get(job_id)   # idempotent re-runs
    except NotFoundError:
        return fw.supervised_fine_tuning_jobs.create(
            supervised_fine_tuning_job_id=job_id,
            base_model=BASE_MODEL,
            dataset=dataset_name,
            output_model=f"accounts/{FIREWORKS_ACCOUNT}/models/{job_id}",
            epochs=2,
            lora_rank=8,
        )

ensure_sft_job("clinical-note-extractor-synthetic", synth_dataset)
ensure_sft_job("clinical-note-extractor-real", real_dataset)

clinical-note-extractor-synthetic: JOB_STATE_COMPLETED
clinical-note-extractor-real:      JOB_STATE_COMPLETED

The jobs run concurrently and finish in roughly 10–20 minutes at this dataset size. Progress, loss curves, and sample renders are visible in the Fireworks dashboard.