{ "cells": [ { "cell_type": "markdown", "id": "313eccc5-d4c7-4ee9-8082-2a1e3670d26d", "metadata": {}, "source": [ "# NuExtract 2.0 Supervised Fine-tuning (SFT)\n", "\n", "This notebook will show a basic example of how to perform supervised fine-tuning (SFT) on top of the base NuExtract 2.0 models, with your own data.\n", "\n", "## Prepare Model\n", "First, load the model you want to fine-tune, along with the processor." ] }, { "cell_type": "code", "execution_count": 1, "id": "cd9abd77-b0d6-4fac-9964-05da8e23f1de", "metadata": { "tags": [] }, "outputs": [], "source": [ "import torch\n", "from transformers import AutoProcessor, AutoModelForVision2Seq\n", "from qwen_vl_utils import process_vision_info\n", "\n", "model_name = \"numind/NuExtract-2.0-2B\"\n", "# model_name = \"numind/NuExtract-2.0-8B\"\n", "\n", "model = AutoModelForVision2Seq.from_pretrained(model_name, \n", " trust_remote_code=True, \n", " torch_dtype=torch.bfloat16,\n", " attn_implementation=\"flash_attention_2\",\n", " device_map=\"auto\",\n", " use_cache=False, # for training\n", " )\n", "\n", "processor = AutoProcessor.from_pretrained(model_name, \n", " trust_remote_code=True, \n", " padding_side='right', # make sure to set padding to right for training\n", " use_fast=True,\n", " )\n", "processor.eos_token = processor.tokenizer.eos_token\n", "processor.eos_token_id = processor.tokenizer.eos_token_id" ] }, { "cell_type": "markdown", "id": "861218a0-9323-495a-b82e-d9fc4bde5670", "metadata": {}, "source": [ "## Prepare Data\n", "\n", "The `construct_messages()` function below will help us to format the messages to be fed into the model." ] }, { "cell_type": "code", "execution_count": 2, "id": "05ad2f4c-116b-48ec-9a6b-e70ec0a3f996", "metadata": { "tags": [] }, "outputs": [], "source": [ "def construct_messages(document, template, label=None, examples=None, image_placeholder=\"<|vision_start|><|image_pad|><|vision_end|>\"):\n", " \"\"\"\n", " Construct the individual NuExtract message texts, prior to chat template formatting.\n", " \"\"\"\n", " images = []\n", " \n", " # add few-shot examples if needed\n", " icl = \"\"\n", " if examples is not None and len(examples) > 0:\n", " icl = \"# Examples:\\n\"\n", " for row in examples:\n", " example_input = row['input']\n", " \n", " if not isinstance(row['input'], str):\n", " example_input = image_placeholder\n", " images.append(row['input'])\n", " \n", " icl += f\"## Input:\\n{example_input}\\n## Output:\\n{row['output']}\\n\"\n", " \n", " # if input document is an image, set text to an image placeholder\n", " text = document\n", " if not isinstance(document, str):\n", " text = image_placeholder\n", " images.append(document)\n", " text = f\"\"\"# Template:\\n{template}\\n{icl}# Context:\\n{text}\"\"\"\n", " \n", " messages = [\n", " {\n", " \"role\": \"user\",\n", " \"content\": [{\"type\": \"text\", \"text\": text}] + images,\n", " }\n", " ]\n", " if label is not None:\n", " messages.append({\n", " \"role\": \"assistant\",\n", " \"content\": [{\"type\": \"text\", \"text\": label}],\n", " \n", " })\n", " return messages" ] }, { "cell_type": "markdown", "id": "ee16acd5-c327-4a6b-b2c9-448a766fa3c0", "metadata": {}, "source": [ "For illustration purposes, we will use a small dataset of manually created examples. You should prepare your own data in a similar way before fine-tuning your own model.\n", "\n", "In the custom data below, we will only provide examples that return strings in full lowercase characters (unless ICL examples suggest otherwise). If we fine-tune on this, we would ideally alter the model to favour returning strings in lowercase by default.\n", "\n", "*Note: training on a very small dataset like this is for illustration purposes only and will almost always result in a poorly performing model in real use-cases.*" ] }, { "cell_type": "code", "execution_count": 3, "id": "4443f697-b11e-4818-b679-b0de2d8eb055", "metadata": { "tags": [] }, "outputs": [], "source": [ "inputs = [\n", " # image input with no ICL examples\n", " {\n", " \"document\": {\"type\": \"image\", \"image\": \"file://data/0.jpg\"},\n", " \"template\": \"\"\"{\"store_name\": \"verbatim-string\"}\"\"\",\n", " \"label\": \"\"\"{\"store_name\": \"walmart\"}\"\"\", # lowercase result\n", " },\n", " # image input with 1 ICL example\n", " {\n", " \"document\": {\"type\": \"image\", \"image\": \"file://data/1.jpg\"},\n", " \"template\": \"\"\"{\"store_name\": \"verbatim-string\"}\"\"\",\n", " \"examples\": [\n", " {\n", " \"input\": {\"type\": \"image\", \"image\": \"file://data/0.jpg\"},\n", " \"output\": \"\"\"{\"store_name\": \"Walmart\"}\"\"\",\n", " }\n", " ],\n", " \"label\": \"\"\"{\"store_name\": \"Trader Joe's\"}\"\"\",\n", " },\n", " # text input with no ICL examples\n", " {\n", " \"document\": \"John went to the restaurant with Mary. James went to the cinema.\",\n", " \"template\": \"\"\"{\"names\": [\"verbatim-string\"]}\"\"\",\n", " \"label\": \"\"\"{\"names\": [\"john\", \"mary\", \"james\"]}\"\"\", # lowercase result\n", " },\n", " # text input with ICL example\n", " {\n", " \"document\": \"John went to the restaurant with Mary. James went to the cinema.\",\n", " \"template\": \"\"\"{\"names\": [\"verbatim-string\"]}\"\"\",\n", " \"examples\": [\n", " {\n", " \"input\": \"Stephen is the manager at Susan's store.\",\n", " \"output\": \"\"\"{\"names\": [\"STEPHEN\", \"SUSAN\"]}\"\"\"\n", " }\n", " ],\n", " \"label\": \"\"\"{\"names\": [\"JOHN\", \"MARY\", \"JAMES\"]}\"\"\",\n", " },\n", "] * 2 # double examples to have dataset of size 8\n", "\n", "messages = [\n", " construct_messages(\n", " x[\"document\"], \n", " x[\"template\"], \n", " x[\"label\"],\n", " x[\"examples\"] if \"examples\" in x else None\n", " ) for x in inputs\n", "]" ] }, { "cell_type": "code", "execution_count": 4, "id": "e4428274-62b1-4fe7-9296-e702b1d9aba5", "metadata": { "tags": [] }, "outputs": [ { "data": { "text/plain": [ "[{'role': 'user',\n", " 'content': [{'type': 'text',\n", " 'text': '# Template:\\n{\"store_name\": \"verbatim-string\"}\\n# Context:\\n<|vision_start|><|image_pad|><|vision_end|>'},\n", " {'type': 'image', 'image': 'file://data/0.jpg'}]},\n", " {'role': 'assistant',\n", " 'content': [{'type': 'text', 'text': '{\"store_name\": \"walmart\"}'}]}]" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "messages[0]" ] }, { "cell_type": "markdown", "id": "751f3ec5-9e3e-4951-b21a-53b8b03209dc", "metadata": {}, "source": [ "Let's also add a couple of validation examples that we can use to confirm our model is generalizing to unseen data." ] }, { "cell_type": "code", "execution_count": 5, "id": "d619853f-d76c-4d28-aebb-0d26b99248fa", "metadata": { "tags": [] }, "outputs": [], "source": [ "val_inputs = [\n", " {\n", " \"document\": \"Jack went to the hill with Jill. Rupert went to the diner.\",\n", " \"template\": \"\"\"{\"names\": [\"verbatim-string\"]}\"\"\",\n", " \"label\": \"\"\"{\"names\": [\"jack\", \"jill\", \"rupert\"]}\"\"\", # lowercase result\n", " },\n", " {\n", " \"document\": \"My dog Clifford likes to play fetch with Emily and Peter.\",\n", " \"template\": \"\"\"{\"names\": [\"verbatim-string\"]}\"\"\",\n", " \"label\": \"\"\"{\"names\": [\"clifford\", \"emily\", \"peter\"]}\"\"\", # lowercase result\n", " },\n", "]\n", "\n", "val_messages = [\n", " construct_messages(\n", " x[\"document\"], \n", " x[\"template\"], \n", " x[\"label\"],\n", " x[\"examples\"] if \"examples\" in x else None\n", " ) for x in val_inputs\n", "]" ] }, { "cell_type": "markdown", "id": "8c129755-5a99-43e6-98d4-02fa8e1e16d1", "metadata": {}, "source": [ "The data is now structued in message format that the processor will be able to reformat via the chat template before tokenization. We will do that on the fly during training via the collate function below." ] }, { "cell_type": "code", "execution_count": 6, "id": "dfac13eb-8ede-4939-b40b-a515c0d73110", "metadata": { "tags": [] }, "outputs": [], "source": [ "def collate_fn(examples):\n", " # process input/prompt part of conversations\n", " user_texts = [processor.apply_chat_template(example[:1], tokenize=False) for example in examples]\n", " \n", " # process full conversations (user + assistant)\n", " full_texts = [processor.apply_chat_template(example, tokenize=False) for example in examples]\n", " \n", " # process images\n", " images = process_vision_info(examples)[0]\n", " \n", " # tokenize sequences\n", " user_batch = processor(text=user_texts, images=images, return_tensors=\"pt\", padding=True)\n", " full_batch = processor(text=full_texts, images=images, return_tensors=\"pt\", padding=True)\n", " \n", " # mask padding tokens\n", " labels = full_batch[\"input_ids\"].clone()\n", " labels[labels == processor.tokenizer.pad_token_id] = -100\n", " \n", " # mask user message tokens for each example in the batch\n", " for i in range(len(examples)):\n", " # length of prompt message (accounting for possible padding)\n", " user_len = user_batch[\"attention_mask\"][i].sum().item()\n", " \n", " # mask prompt part of label\n", " labels[i, :user_len - 1] = -100\n", " \n", " full_batch[\"labels\"] = labels\n", " return full_batch" ] }, { "cell_type": "code", "execution_count": 7, "id": "a3311e1d-8d53-4a0f-bc5a-4197e316ae1d", "metadata": { "tags": [] }, "outputs": [ { "data": { "text/plain": [ "torch.Size([2688, 1176])" ] }, "execution_count": 7, "metadata": {}, "output_type": "execute_result" } ], "source": [ "collate_fn(messages[:1])['pixel_values'].shape" ] }, { "cell_type": "markdown", "id": "6139f83c-d83c-4b4b-8e87-e28320799b5c", "metadata": {}, "source": [ "## Fine-Tune the Model\n", "\n", "We will use the `SFTTrainer` from the `trl` library, which abstracts a lot of the complexities of training for us. For your own use-case you should adjust various hyper-parameters like learning rate, epochs, etc. according to your problem." ] }, { "cell_type": "code", "execution_count": 8, "id": "8df597dd-7eac-4331-ad2c-3fc6800ea4e3", "metadata": { "tags": [] }, "outputs": [], "source": [ "from trl import SFTConfig, SFTTrainer\n", "\n", "# Configure training arguments\n", "training_args = SFTConfig(\n", " output_dir=\"test_finetune\", # Directory to save the model\n", " num_train_epochs=5, # Number of training epochs\n", " per_device_train_batch_size=1, # Batch size for training\n", " per_device_eval_batch_size=1, # Batch size for evaluation\n", " gradient_accumulation_steps=4, # Steps to accumulate gradients\n", " learning_rate=1e-5, # Learning rate for training\n", " lr_scheduler_type=\"constant\", # Type of learning rate scheduler\n", " logging_steps=1, # Steps interval for logging\n", " eval_steps=2, # Steps interval for evaluation\n", " eval_strategy=\"steps\", # Strategy for evaluation\n", " # save_strategy=\"steps\", # Strategy for saving the model\n", " # save_steps=20, # Steps interval for saving\n", " bf16=True, # Use bfloat16 precision\n", " max_grad_norm=0.3, # Maximum norm for gradient clipping\n", " warmup_ratio=0.03, # Ratio of total steps for warmup\n", " report_to=\"none\", # Reporting tool for tracking metrics\n", " gradient_checkpointing=True, # Enable gradient checkpointing for memory efficiency\n", " gradient_checkpointing_kwargs={\"use_reentrant\": False}, # Options for gradient checkpointing\n", " # max_seq_length=1024 # Maximum sequence length for input\n", ")\n", "\n", "# allow for proper loading of images during collation\n", "training_args.remove_unused_columns = False\n", "training_args.dataset_kwargs = {\"skip_prepare_dataset\": True}\n", "\n", "trainer = SFTTrainer(\n", " model=model,\n", " args=training_args,\n", " data_collator=collate_fn,\n", " train_dataset=messages,\n", " eval_dataset=val_messages,\n", " processing_class=processor.tokenizer,\n", ")" ] }, { "cell_type": "code", "execution_count": 9, "id": "0d05bdce-dc6e-4979-a573-8aed27972291", "metadata": { "tags": [] }, "outputs": [ { "data": { "text/html": [ "\n", "
\n", " \n", " \n", " [10/10 01:34, Epoch 5/5]\n", "
\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
StepTraining LossValidation Loss
20.1704000.269275
40.0183000.042456
60.0018000.012193
80.0003000.011337
100.0003000.011667

" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "TrainOutput(global_step=10, training_loss=0.057332569236314156, metrics={'train_runtime': 96.9988, 'train_samples_per_second': 0.412, 'train_steps_per_second': 0.103, 'total_flos': 261136381470720.0, 'train_loss': 0.057332569236314156})" ] }, "execution_count": 9, "metadata": {}, "output_type": "execute_result" } ], "source": [ "trainer.train()" ] }, { "cell_type": "code", "execution_count": 10, "id": "7c800e94-9494-4178-b987-d1eebbd80ec0", "metadata": { "tags": [] }, "outputs": [], "source": [ "trainer.save_model(training_args.output_dir)" ] }, { "cell_type": "markdown", "id": "1c8cb789-9168-4b68-ae57-2d67ab104743", "metadata": { "tags": [] }, "source": [ "## Test Generation\n", "\n", "Now, let's run actual generation of outputs for our validation examples to see if what the model has learned." ] }, { "cell_type": "code", "execution_count": 11, "id": "a66eb36d-8301-4e43-8bcd-8dc831a625b7", "metadata": { "tags": [] }, "outputs": [], "source": [ "# reload processor with left padding (for generation)\n", "processor = AutoProcessor.from_pretrained(model_name, \n", " trust_remote_code=True, \n", " padding_side='left',\n", " use_fast=True)\n", "\n", "# reconstruct validation messages without labels\n", "test_messages = [\n", " construct_messages(\n", " x[\"document\"], \n", " x[\"template\"], \n", " ) for x in val_inputs\n", "]\n", "\n", "texts = processor.tokenizer.apply_chat_template(\n", " test_messages,\n", " tokenize=False,\n", " add_generation_prompt=True,\n", ")\n", "\n", "image_inputs = process_vision_info(messages[2][:1])[0]\n", "inputs = processor(\n", " text=texts,\n", " images=image_inputs,\n", " padding=True,\n", " return_tensors=\"pt\",\n", ").to(\"cuda\")\n", "\n", "# we choose greedy sampling here, which works well for most information extraction tasks\n", "generation_config = {\"do_sample\": True, \"temperature\": 1.0, \"max_new_tokens\": 2048}\n", "\n", "# Inference: Generation of the output\n", "generated_ids = model.generate(\n", " **inputs,\n", " **generation_config\n", ")\n", "generated_ids_trimmed = [\n", " out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)\n", "]\n", "output_texts = processor.batch_decode(\n", " generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False\n", ")" ] }, { "cell_type": "markdown", "id": "14609ef3-7121-4956-b25e-6ded1a12bc70", "metadata": {}, "source": [ "As can be seen below, the model is now generating extractions with lowercase strings by default." ] }, { "cell_type": "code", "execution_count": 12, "id": "091058ee-347f-45dc-aed8-21693a818334", "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "=== Prompt ===\n", "<|im_start|>system\n", "You are a helpful assistant.<|im_end|>\n", "<|im_start|>user\n", "# Template:\n", "{\"names\": [\"verbatim-string\"]}\n", "# Context:\n", "Jack went to the hill with Jill. Rupert went to the diner.<|im_end|>\n", "<|im_start|>assistant\n", "\n", "=== Output ===\n", "{\"names\": [\"jack\", \"jill\", \"rupert\"]}\n", "\n", "=== Prompt ===\n", "<|im_start|>system\n", "You are a helpful assistant.<|im_end|>\n", "<|im_start|>user\n", "# Template:\n", "{\"names\": [\"verbatim-string\"]}\n", "# Context:\n", "My dog Clifford likes to play fetch with Emily and Peter.<|im_end|>\n", "<|im_start|>assistant\n", "\n", "=== Output ===\n", "{\"names\": [\"clifford\", \"emily\", \"peter\"]}\n", "\n" ] } ], "source": [ "for i in range(len(texts)):\n", " print(f\"=== Prompt ===\\n{texts[i]}\")\n", " print(f\"=== Output ===\\n{output_texts[i]}\\n\")" ] }, { "cell_type": "code", "execution_count": null, "id": "eb987296-d418-4cbb-b97c-37f4dd27632b", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "environment": { "kernel": "conda-base-py", "name": "workbench-notebooks.m127", "type": "gcloud", "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m127" }, "kernelspec": { "display_name": "Python 3 (ipykernel) (Local)", "language": "python", "name": "conda-base-py" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.16" } }, "nbformat": 4, "nbformat_minor": 5 }