{ "cells": [ { "cell_type": "markdown", "id": "c88a5f45", "metadata": { "id": "c88a5f45" }, "source": [ "# Fine-Tuning LFM3-1.7B for Structured Extraction with TRL (SFT)\n", "\n", "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Liquid4All/cookbook/blob/main/finetuning/notebooks/lfm3_sft_with_trl.ipynb) [![GitHub](https://img.shields.io/badge/GitHub-Liquid4All%2Fcookbook-181717?logo=github)](https://github.com/Liquid4All/cookbook) [![Discord](https://img.shields.io/discord/1385439864920739850?color=7289da&label=Discord&logo=discord&logoColor=white)](https://discord.com/invite/liquid-ai)\n", "\n", "This notebook fine-tunes [LFM3-1.7B](https://huggingface.co/LiquidAI/LFM3-1.7B) to turn messy real-world documents (emails, receipts, booking confirmations, forms in many languages) into clean JSON that follows a schema. We use Supervised Fine-Tuning (SFT) with [TRL](https://github.com/huggingface/trl) and LoRA, and measure field-level accuracy before and after training, including on document types the model never saw during training.\n", "\n", "Extraction is a typical on-device task: the documents stay on the device, and a small model that reliably fills your schema is faster and cheaper than a large general one.\n", "\n", "## šŸŽÆ What you'll find\n", "\n", "- **The task**: a JSON schema and a document go in, a JSON object comes out\n", "- **LoRA + SFT**: train on the JSON answers only, in LFM3's non-reasoning mode\n", "- **Evaluation**: valid JSON rate, field accuracy, and exact match, on seen and unseen document types" ] }, { "cell_type": "markdown", "id": "4bc56ae2", "metadata": { "id": "4bc56ae2" }, "source": [ "# šŸ“¦ Installation & Setup\n", "\n", "We pin every library to the versions this notebook was tested with." ] }, { "cell_type": "code", "execution_count": null, "id": "fd0cc5c1", "metadata": { "tags": [ "install" ], "id": "fd0cc5c1" }, "outputs": [], "source": [ "!pip install -qqq \"torch==2.14.1\" \"torchvision==0.29.1\" \"torchao==0.18.0\" \"trl==1.16.0\" \"transformers==5.20.2\" \"peft==0.21.2\" \"accelerate==1.15.0\" \"datasets==5.1.0\" \"huggingface-hub==2.3.0\" \"tokenizers==0.23.3\" \"safetensors==0.8.0\" --progress-bar off" ] }, { "cell_type": "code", "execution_count": null, "id": "b6a7fc3c", "metadata": { "id": "b6a7fc3c" }, "outputs": [], "source": [ "import json\n", "import os\n", "import torch\n", "import transformers\n", "import trl\n", "import peft\n", "\n", "os.environ[\"WANDB_DISABLED\"] = \"true\"\n", "\n", "print(f\"šŸ“¦ PyTorch version: {torch.__version__}\")\n", "print(f\"šŸ¤— Transformers version: {transformers.__version__}\")\n", "print(f\"šŸ“Š TRL version: {trl.__version__}\")\n", "print(f\"šŸŽ›ļø PEFT version: {peft.__version__}\")\n", "\n", "# LFM3-1.7B is released in bfloat16. GPUs without bfloat16 support, like the T4, load it in float16 instead:\n", "# the conversion is exact for almost every weight, and the LoRA adapters train in float32 either way.\n", "use_bf16 = torch.cuda.is_bf16_supported(including_emulation=False)\n", "dtype = torch.bfloat16 if use_bf16 else torch.float16\n", "print(f\"šŸ”¢ Precision: {dtype}\")" ] }, { "cell_type": "markdown", "id": "b07c3d9b", "metadata": { "id": "b07c3d9b" }, "source": [ "# Loading the model from Transformers šŸ¤—\n", "\n", "LFM3-1.7B ships its modeling code with the checkpoint, so we load it with `trust_remote_code=True`. We pad on the left, so that batched generation continues each prompt right where it ends." ] }, { "cell_type": "code", "execution_count": null, "id": "55eb8430", "metadata": { "id": "55eb8430" }, "outputs": [], "source": [ "from transformers import AutoModelForCausalLM, AutoTokenizer\n", "\n", "model_id = \"LiquidAI/LFM3-1.7B\"\n", "\n", "print(\"šŸ“š Loading tokenizer...\")\n", "tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True, padding_side=\"left\")\n", "\n", "print(\"🧠 Loading model...\")\n", "model = AutoModelForCausalLM.from_pretrained(model_id, dtype=dtype, device_map=\"auto\", trust_remote_code=True)\n", "\n", "print(\"āœ… Model loaded successfully!\")\n", "print(f\"šŸ”¢ Parameters: {model.num_parameters():,}\")" ] }, { "cell_type": "markdown", "id": "a85fa0e1", "metadata": { "id": "a85fa0e1" }, "source": [ "# šŸ“š Load the extraction dataset\n", "\n", "Each example has a document type, its JSON schema, a document, and the gold extraction. The documents are fictional but realistic, in 16 languages, and some leave fields unanswered on purpose: the gold value is then `null`, and the model must not invent one.\n", "\n", "The dataset has three splits: `train`, `test` (new documents of the training document types), and `test_unseen` (document types that never appear in training)." ] }, { "cell_type": "code", "execution_count": null, "id": "7c8f0a51", "metadata": { "id": "7c8f0a51" }, "outputs": [], "source": [ "from datasets import load_dataset\n", "\n", "dataset = load_dataset(\"LiquidAI/structured-extraction\")\n", "print(dataset)\n", "\n", "example = dataset[\"train\"][0]\n", "print(f\"\\nšŸ“„ Document type: {example['doc_type']} ({example['language']})\")\n", "print(f\"\\nšŸ“ Document:\\n{example['document'][:800]}\")\n", "print(f\"\\nšŸŽÆ Gold extraction:\\n{json.dumps(json.loads(example['extraction']), indent=2, ensure_ascii=False)[:800]}\")" ] }, { "cell_type": "markdown", "id": "e72671f0", "metadata": { "id": "e72671f0" }, "source": [ "## Build the prompts\n", "\n", "The system prompt holds the extraction rules the dataset follows (date and number formats, `null` for missing information, ...) and the schema. The document goes in the user message, and the answer is the JSON object, on one line. Both the original and the fine-tuned model get the same instructions.\n", "\n", "Extraction doesn't need step-by-step reasoning, so we use LFM3's non-reasoning mode (`enable_thinking=False`): the answer comes right away, which is faster and cheaper. TRL reads the `chat_template_kwargs` column and passes it to the chat template." ] }, { "cell_type": "code", "execution_count": null, "id": "84370dc3", "metadata": { "id": "84370dc3" }, "outputs": [], "source": [ "SYSTEM_PROMPT = \"\"\"Extract the information from the document into a JSON object that follows the schema below, with these rules:\n", "\n", "1. Output exactly one JSON object with exactly the keys of the schema, in schema order, at every nesting level. No extra keys, no prose, no code fences.\n", "2. String values are copied from the document exactly as written: same language, script, spelling and capitalization. Never translate or paraphrase. Leave out field labels (\"Order no.:\"), surrounding quotes and trailing punctuation. For person names, leave out honorifics and courtesy titles (Mr, Mrs, Dr, Prof, Herr, Frau, Mme, M., Sr., Sra., ꧘, 恕悓, å…ˆē”Ÿ, ė‹˜, г-н). Phone numbers, IDs, emails, URLs and codes are copied as written. A value printed over several lines is joined with \", \".\n", "3. Dates are `YYYY-MM-DD` and times are `HH:MM` on a 24-hour clock, whatever the document's local format, language or calendar (a Buddhist-era year 2569 is 2026).\n", "4. Amounts and quantities are JSON numbers: no currency symbol, no thousands separator, dot as the decimal separator (`1.234,50 €` -> `1234.5`, `Ā„12,800` -> `12800`).\n", "5. A currency field holds the ISO 4217 code (EUR, USD, JPY, ...) of the amounts in the document.\n", "6. Booleans are `true` or `false` only when the document states it; otherwise `null`.\n", "7. Enum fields take one of the schema's allowed values, chosen by meaning even when the document is in another language.\n", "8. If the document does not contain the information for a field, the value is `null`. Never guess, infer from outside knowledge, or compute a value that is not printed. An array with no entries is `[]`.\n", "9. Array items follow the order they appear in the document.\n", "10. When the document mentions several candidate values (an old and a new date, a previous balance and the amount due, a CC'd person and an attendee), use the one the field description asks for; updated values replace superseded ones.\n", "\n", "Schema:\n", "{schema}\"\"\"\n", "\n", "def build_messages(example):\n", " return [\n", " {\"role\": \"system\", \"content\": SYSTEM_PROMPT.format(schema=example[\"schema\"])},\n", " {\"role\": \"user\", \"content\": example[\"document\"]},\n", " ]\n", "\n", "def to_prompt_completion(example):\n", " answer = json.dumps(json.loads(example[\"extraction\"]), ensure_ascii=False)\n", " return {\n", " \"prompt\": build_messages(example),\n", " \"completion\": [{\"role\": \"assistant\", \"content\": answer}],\n", " \"chat_template_kwargs\": {\"enable_thinking\": False},\n", " }\n", "\n", "train_dataset = dataset[\"train\"].map(to_prompt_completion, remove_columns=dataset[\"train\"].column_names)\n", "print(tokenizer.apply_chat_template(train_dataset[0][\"prompt\"] + train_dataset[0][\"completion\"], tokenize=False, enable_thinking=False)[-700:])" ] }, { "cell_type": "markdown", "id": "4a73b586", "metadata": { "id": "4a73b586" }, "source": [ "# šŸ“ Evaluate the model before training\n", "\n", "We generate an extraction for each evaluation document, with greedy decoding and the repetition penalty recommended in the model card (`repetition_penalty=1.1`), and compare it with the gold JSON:\n", "\n", "- **Valid JSON**: the output parses as a JSON object (we tolerate a Markdown code fence around it, even though the rules ask for none).\n", "- **Field accuracy**: the share of fields with the right value. Strings are compared without case and extra spaces, numbers as numbers, and `null` must stay `null`.\n", "- **Exact match**: every field is right.\n", "\n", "To keep the notebook fast, we evaluate on 200 documents of each test split." ] }, { "cell_type": "code", "execution_count": null, "id": "09c98400", "metadata": { "id": "09c98400" }, "outputs": [], "source": [ "from tqdm.auto import tqdm\n", "\n", "def parse_json(text):\n", " \"\"\"The JSON object in a model's answer, tolerating a Markdown code fence around it, or None.\"\"\"\n", " text = text.strip()\n", " if text.startswith(\"```\"):\n", " text = text.split(\"\\n\", 1)[-1].rsplit(\"```\", 1)[0]\n", " try:\n", " value = json.loads(text)\n", " except json.JSONDecodeError:\n", " return None\n", " return value if isinstance(value, dict) else None\n", "\n", "def flatten(value, path=\"\"):\n", " \"\"\"Leaves of a JSON value as {path: value}, e.g. {\"items[0].price\": 4.5}.\"\"\"\n", " if isinstance(value, dict):\n", " return {k: v for key, child in value.items() for k, v in flatten(child, f\"{path}.{key}\" if path else key).items()}\n", " if isinstance(value, list):\n", " return {k: v for i, child in enumerate(value) for k, v in flatten(child, f\"{path}[{i}]\").items()} or {path: []}\n", " return {path: value}\n", "\n", "def same(a, b):\n", " if isinstance(a, bool) or isinstance(b, bool):\n", " return a is b\n", " if isinstance(a, (int, float)) and isinstance(b, (int, float)):\n", " return abs(a - b) < 1e-6\n", " if isinstance(a, str) and isinstance(b, str):\n", " return \" \".join(a.split()).casefold() == \" \".join(b.split()).casefold()\n", " return a == b\n", "\n", "def score(prediction, gold):\n", " \"\"\"Field accuracy of one prediction: matching fields / fields in the gold or the prediction.\"\"\"\n", " gold_fields, predicted_fields = flatten(gold), flatten(prediction)\n", " paths = set(gold_fields) | set(predicted_fields)\n", " if not paths:\n", " return 1.0\n", " return sum(p in gold_fields and p in predicted_fields and same(gold_fields[p], predicted_fields[p]) for p in paths) / len(paths)\n", "\n", "@torch.no_grad()\n", "def evaluate(model, split, n=200, batch_size=16):\n", " model.eval()\n", " examples = dataset[split].shuffle(seed=42).select(range(n))\n", " valid, accuracy, exact = 0, 0.0, 0\n", " for start in tqdm(range(0, n, batch_size), desc=f\"Evaluating {split}\"):\n", " batch = examples.select(range(start, min(start + batch_size, n)))\n", " prompts = [tokenizer.apply_chat_template(build_messages(e), add_generation_prompt=True, tokenize=False, enable_thinking=False) for e in batch]\n", " inputs = tokenizer(prompts, return_tensors=\"pt\", padding=True, add_special_tokens=False).to(model.device)\n", " outputs = model.generate(**inputs, max_new_tokens=1024, do_sample=False, repetition_penalty=1.1)\n", " for output, example in zip(outputs[:, inputs[\"input_ids\"].shape[1]:], batch):\n", " gold = json.loads(example[\"extraction\"])\n", " prediction = parse_json(tokenizer.decode(output, skip_special_tokens=True))\n", " if prediction is None:\n", " continue # invalid JSON scores 0 on every metric\n", " valid += 1\n", " field_accuracy = score(prediction, gold)\n", " accuracy += field_accuracy\n", " exact += field_accuracy == 1.0\n", " return {\"valid_json\": valid / n, \"field_accuracy\": accuracy / n, \"exact_match\": exact / n}\n", "\n", "def show(results, title):\n", " print(title)\n", " for split, r in results.items():\n", " print(f\" {split:<12} valid JSON {r['valid_json']:.1%} | field accuracy {r['field_accuracy']:.1%} | exact match {r['exact_match']:.1%}\")\n", "\n", "baseline = {split: evaluate(model, split) for split in [\"test\", \"test_unseen\"]}\n", "show(baseline, \"šŸ“ Before training\")" ] }, { "cell_type": "markdown", "id": "2e45d0ae", "metadata": { "id": "2e45d0ae" }, "source": [ "# šŸŽÆ Supervised Fine-Tuning (SFT + LoRA)\n", "\n", "## Wrap the model with PEFT\n", "\n", "We adapt the attention projections (`q_proj`, `k_proj`, `v_proj`, `o_proj`), the gated short convolutions (`in_proj`, `out_proj`), and the MLP (`w1`, `w2`, `w3`). LFM3's attention layers also have a per-head output gate, `gate_proj`; adding it to the targets is a reasonable choice." ] }, { "cell_type": "code", "execution_count": null, "id": "96cd6067", "metadata": { "id": "96cd6067" }, "outputs": [], "source": [ "from peft import LoraConfig\n", "\n", "peft_config = LoraConfig(\n", " r=16,\n", " lora_alpha=32,\n", " lora_dropout=0.05,\n", " target_modules=[\n", " \"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"in_proj\", \"out_proj\", \"w1\", \"w2\", \"w3\",\n", " # \"gate_proj\", # optional: LFM3's per-head attention gate\n", " ],\n", " task_type=\"CAUSAL_LM\",\n", ")" ] }, { "cell_type": "markdown", "id": "2bb226b0", "metadata": { "id": "2bb226b0" }, "source": [ "## Launch training\n", "\n", "With a prompt-completion dataset, `SFTTrainer` computes the loss on the completion only: the model learns to write the JSON, not to reproduce the schema or the document." ] }, { "cell_type": "code", "execution_count": null, "id": "452c410e", "metadata": { "id": "452c410e" }, "outputs": [], "source": [ "from trl import SFTConfig, SFTTrainer\n", "\n", "sft_config = SFTConfig(\n", " output_dir=\"./LFM3-1.7B-extraction\",\n", " num_train_epochs=1,\n", " per_device_train_batch_size=4,\n", " gradient_accumulation_steps=4, # <- effective batch size of 16 documents\n", " learning_rate=2e-4,\n", " lr_scheduler_type=\"cosine\",\n", " warmup_steps=0.05, # <- a float is a fraction of the total steps\n", " max_length=4096,\n", " logging_steps=10,\n", " save_strategy=\"no\",\n", " bf16=use_bf16,\n", " fp16=not use_bf16, # <- float16 mixed precision with loss scaling on GPUs without bfloat16\n", " report_to=\"none\",\n", ")\n", "\n", "trainer = SFTTrainer(\n", " model=model,\n", " args=sft_config,\n", " train_dataset=train_dataset,\n", " processing_class=tokenizer,\n", " peft_config=peft_config,\n", ")\n", "trainer.model.print_trainable_parameters()\n", "\n", "print(\"šŸš€ Starting SFT training...\")\n", "trainer.train()\n", "print(\"šŸŽ‰ SFT training completed!\")\n", "print(f\"🧮 Peak GPU memory: {torch.cuda.max_memory_reserved() / 1e9:.1f} GB\")" ] }, { "cell_type": "markdown", "id": "ab45b4e1", "metadata": { "id": "ab45b4e1" }, "source": [ "# šŸ“Š Evaluate the fine-tuned model" ] }, { "cell_type": "code", "execution_count": null, "id": "0135b0f7", "metadata": { "id": "0135b0f7" }, "outputs": [], "source": [ "finetuned = {split: evaluate(trainer.model, split) for split in [\"test\", \"test_unseen\"]}\n", "show(baseline, \"šŸ“ Before training\")\n", "show(finetuned, \"šŸ“ After training\")" ] }, { "cell_type": "markdown", "id": "d5cd1aba", "metadata": { "id": "d5cd1aba" }, "source": [ "# šŸ” Original vs fine-tuned model\n", "\n", "To see what changed, we run the original model and the fine-tuned model on the same documents. With LoRA, we don't need to load the original model again: `disable_adapter()` temporarily turns the adapter off.\n", "\n", "The `showcase` split holds a few documents of types the model never saw during training. For each field, we show the gold value and both extractions, with āœ… when the value is right." ] }, { "cell_type": "code", "execution_count": null, "id": "fac9de60", "metadata": { "id": "fac9de60" }, "outputs": [], "source": [ "def extract(model, example):\n", " prompt = tokenizer.apply_chat_template(build_messages(example), add_generation_prompt=True, tokenize=False, enable_thinking=False)\n", " inputs = tokenizer(prompt, return_tensors=\"pt\", add_special_tokens=False).to(model.device)\n", " output = model.generate(**inputs, max_new_tokens=1024, do_sample=False, repetition_penalty=1.1)\n", " return parse_json(tokenizer.decode(output[0, inputs[\"input_ids\"].shape[1]:], skip_special_tokens=True))\n", "\n", "def cell(fields, path, gold_fields):\n", " if fields is None:\n", " return \"āŒ invalid JSON\"\n", " if path not in fields:\n", " return \"āŒ (missing)\"\n", " mark = \"āœ…\" if path in gold_fields and same(fields[path], gold_fields[path]) else \"āŒ\"\n", " return f\"{mark} {json.dumps(fields[path], ensure_ascii=False)}\"\n", "\n", "trainer.model.eval()\n", "for example in dataset[\"showcase\"]:\n", " with trainer.model.disable_adapter():\n", " original = extract(trainer.model, example)\n", " finetuned_extraction = extract(trainer.model, example)\n", " gold_fields = flatten(json.loads(example[\"extraction\"]))\n", " original_fields = flatten(original) if isinstance(original, dict) else None\n", " finetuned_fields = flatten(finetuned_extraction) if isinstance(finetuned_extraction, dict) else None\n", " print(f\"\\nšŸ“„ {example['doc_type']} ({example['language']})\\n{example['document'][:400]}...\\n\")\n", " print(f\"{'field':<28} {'gold':<30} {'original':<34} {'fine-tuned'}\")\n", " for path, gold_value in gold_fields.items():\n", " print(f\"{path[:27]:<28} {json.dumps(gold_value, ensure_ascii=False)[:29]:<30} {cell(original_fields, path, gold_fields)[:33]:<34} {cell(finetuned_fields, path, gold_fields)[:40]}\")" ] }, { "cell_type": "markdown", "id": "78fb8cb0", "metadata": { "id": "78fb8cb0" }, "source": [ "And the summary over all the evaluation documents:" ] }, { "cell_type": "code", "execution_count": null, "id": "001e9edd", "metadata": { "id": "001e9edd" }, "outputs": [], "source": [ "print(f\"{'':<34}{'original':>10}{'fine-tuned':>12}\")\n", "for split, name in [(\"test\", \"seen document types\"), (\"test_unseen\", \"unseen document types\")]:\n", " for metric in [\"valid_json\", \"field_accuracy\", \"exact_match\"]:\n", " print(f\"{name + ', ' + metric:<34}{baseline[split][metric]:>10.1%}{finetuned[split][metric]:>12.1%}\")" ] }, { "cell_type": "markdown", "id": "f72ce1f1", "metadata": { "id": "f72ce1f1" }, "source": [ "# šŸ’¾ Save and merge the model\n", "\n", "We save the LoRA adapter, then merge it into the base weights. `save_pretrained` also copies LFM3's modeling code next to the weights, so the merged model loads with `trust_remote_code=True` like the original one." ] }, { "cell_type": "code", "execution_count": null, "id": "35df9090", "metadata": { "id": "35df9090" }, "outputs": [], "source": [ "trainer.save_model(\"./LFM3-1.7B-extraction-lora\")\n", "\n", "output_dir = \"./LFM3-1.7B-extraction-merged\"\n", "merged_model = trainer.model.merge_and_unload()\n", "merged_model.save_pretrained(output_dir)\n", "tokenizer.save_pretrained(output_dir)\n", "print(f\"šŸ’¾ Merged model saved to: {output_dir}\")\n", "\n", "# To push the merged model to the Hugging Face Hub, uncomment:\n", "# merged_model.push_to_hub(\"your-username/LFM3-1.7B-extraction\")\n", "# tokenizer.push_to_hub(\"your-username/LFM3-1.7B-extraction\")" ] }, { "cell_type": "markdown", "id": "369f7a43", "metadata": { "id": "369f7a43" }, "source": [ "# šŸš€ Next steps\n", "\n", "- **Your own schemas**: write a JSON schema per document type and a few hundred labeled documents. Include documents with missing fields, so the model learns to return `null` instead of guessing.\n", "- **Constrained decoding**: at inference, a JSON-schema grammar (supported by llama.cpp and vLLM) guarantees valid JSON on top of the fine-tuned accuracy." ] } ], "metadata": { "accelerator": "GPU", "colab": { "gpuType": "T4", "provenance": [] }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }