{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 💧 LFM3-Omni - Spoken medical image Q&A with Transformers\n",
    "\n",
    "This tutorial demonstrates how to fine-tune our omni model [`LiquidAI/LFM3-Omni-2B`](https://huggingface.co/LiquidAI/LFM3-Omni-2B)\n",
    "on a task that combines **vision and speech**: a clinician looks at a medical image and **asks a question out loud**,\n",
    "and the model picks the right answer.\n",
    "\n",
    "We reuse the medical VQA dataset from our [LFM2.5-VL fine-tuning notebook](https://github.com/Liquid4All/cookbook/blob/main/finetuning/notebooks/sft_for_vision_language_model_with_trl.ipynb),\n",
    "[simwit/omni-med-vqa-mini](https://huggingface.co/datasets/simwit/omni-med-vqa-mini) (CT, MRI, X-ray, ultrasound, dermoscopy,\n",
    "fundus, microscopy and OCT images, with multiple-choice questions), and turn every question into speech with LFM3-Omni's own text-to-speech.\n",
    "\n",
    "## 🎯 What You'll Find:\n",
    "- **Data preparation** - synthesize spoken questions and build image + audio + text conversations\n",
    "- **LoRA + SFT** - parameter-efficient fine-tuning with the 🤗 `Trainer`\n",
    "- **Evaluation** - multiple-choice accuracy before and after fine-tuning\n",
    "\n",
    "## 📋 Prerequisites:\n",
    "- **GPU Runtime**: Select GPU in `Runtime` → `Change runtime type` (an A100 is recommended)\n",
    "- **Hugging Face Account**: For accessing models and datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 📦 Installation & Setup\n",
    "\n",
    "First, let's install all the required packages:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "!pip install -qqq \"transformers==5.20.2\" \"torchao==0.18.0\" accelerate datasets peft pillow librosa soundfile --progress-bar off"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now verify the packages are installed correctly"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "import os\n",
    "import torch\n",
    "import transformers\n",
    "import peft\n",
    "os.environ[\"WANDB_DISABLED\"] = \"true\"\n",
    "\n",
    "print(f\"📦 PyTorch version: {torch.__version__}\")\n",
    "print(f\"🤗 Transformers version: {transformers.__version__}\")\n",
    "print(f\"🧩 PEFT version: {peft.__version__}\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loading the model from Transformers 🤗"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from transformers import AutoModelForMultimodalLM, AutoProcessor\n",
    "\n",
    "model_id = \"LiquidAI/LFM3-Omni-2B\"\n",
    "\n",
    "print(\"📚 Loading processor...\")\n",
    "processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)\n",
    "\n",
    "print(\"🧠 Loading model...\")\n",
    "model = AutoModelForMultimodalLM.from_pretrained(\n",
    "    model_id, dtype=torch.bfloat16, device_map=\"cuda\", trust_remote_code=True\n",
    ")\n",
    "codec = processor.audio_tokenizer.to(model.device)\n",
    "\n",
    "print(\"\\n✅ Model loaded successfully!\")\n",
    "print(f\"🔢 Parameters: {model.num_parameters():,}\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 📚 Load the dataset\n",
    "\n",
    "Each sample has an image, a question, up to four options and the letter of the correct one.\n",
    "We hold out 200 samples for evaluation."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "raw_ds = load_dataset(\"simwit/omni-med-vqa-mini\", split=\"test\").shuffle(seed=42)\n",
    "split = raw_ds.train_test_split(test_size=200, seed=42)\n",
    "train_raw, eval_raw = split[\"train\"], split[\"test\"]\n",
    "\n",
    "print(\"✅ Dataset loaded:\")\n",
    "print(f\"   📚 Train samples: {len(train_raw)}\")\n",
    "print(f\"   🧪 Eval samples: {len(eval_raw)}\")\n",
    "\n",
    "sample = train_raw[0]\n",
    "display(sample[\"image\"].convert(\"RGB\").resize((224, 224)))\n",
    "print(f\"[{sample['modality']}] {sample['input']}\\n✅ Answer: {sample['label']}\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🗣️ Turn the questions into speech\n",
    "\n",
    "The dataset only has written questions, so we let LFM3-Omni read them aloud, alternating a female and a male voice.\n",
    "Many questions repeat (1,146 unique questions for 2,000 samples), so we synthesize each one once.\n",
    "This takes ~10-15 minutes on an A100."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "import librosa\n",
    "import numpy as np\n",
    "from IPython.display import Audio\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "\n",
    "@torch.inference_mode()\n",
    "def speak(text, voice):\n",
    "    inputs = processor.apply_chat_template(\n",
    "        [\n",
    "            {\"role\": \"system\", \"content\": f\"Read this with a {voice} voice in American English.\"},\n",
    "            {\"role\": \"user\", \"content\": text},\n",
    "        ],\n",
    "        add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors=\"pt\",\n",
    "    ).to(model.device)\n",
    "    out = model.generate(**inputs, max_new_tokens=400, generation_mode=\"audio\", audio_temperature=0.4, audio_repetition_penalty=1.2)\n",
    "    wav = codec.decode(out.audio_codes[:1, :, : out.audio_lengths[0]]).audio_values[0].float().cpu().numpy().squeeze()\n",
    "    return librosa.resample(wav, orig_sr=24_000, target_sr=16_000)  # the audio encoder reads 16 kHz\n",
    "\n",
    "\n",
    "torch.manual_seed(0)\n",
    "questions = sorted(set(train_raw[\"question\"]) | set(eval_raw[\"question\"]))\n",
    "spoken = {q: speak(q, \"female\" if i % 2 else \"male\") for i, q in enumerate(tqdm(questions, desc=\"Synthesizing questions\"))}\n",
    "\n",
    "print(f\"🗣️ {len(spoken)} spoken questions, e.g. {sample['question']!r}\")\n",
    "display(Audio(spoken[sample[\"question\"]], rate=16_000))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Format the conversations\n",
    "\n",
    "The image and the spoken question go in the user turn, followed by the written options. The target is the option letter."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "SYSTEM_PROMPT = \"Answer questions about the provided image. Base every answer only on what can be seen. Respond with text.\"\n",
    "\n",
    "\n",
    "def format_sample(sample, with_answer=True):\n",
    "    options = \"\\n\".join(f\"({k}) {sample['option_' + k]}\" for k in \"ABCD\" if sample[\"option_\" + k])\n",
    "    conversation = [\n",
    "        {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
    "        {\"role\": \"user\", \"content\": [\n",
    "            {\"type\": \"image\", \"image\": sample[\"image\"].convert(\"RGB\")},\n",
    "            {\"type\": \"audio\", \"audio\": spoken[sample[\"question\"]]},\n",
    "            {\"type\": \"text\", \"text\": f\"Options:\\n{options}\\nAnswer with the letter of the correct option.\"},\n",
    "        ]},\n",
    "    ]\n",
    "    if with_answer:\n",
    "        conversation.append({\"role\": \"assistant\", \"content\": sample[\"label\"]})\n",
    "    return conversation\n",
    "\n",
    "\n",
    "train_dataset = [format_sample(s) for s in train_raw]\n",
    "print(processor.apply_chat_template(train_dataset[0], tokenize=False))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Collate function\n",
    "\n",
    "The processor renders a batch of conversations, loads the images and audio, and expands each `<|image|>` and `<|audio|>`\n",
    "placeholder to its embeddings. We train only on the answer: every token up to and including the assistant header\n",
    "(`<|im_start|>assistant\\n<think></think>`) is masked with `-100`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "ASSISTANT_HEADER = processor.tokenizer(\"<|im_start|>assistant\\n<think></think>\", add_special_tokens=False)[\"input_ids\"]\n",
    "\n",
    "\n",
    "def answer_start(ids):\n",
    "    n = len(ASSISTANT_HEADER)\n",
    "    return max(i for i in range(len(ids) - n + 1) if ids[i:i + n] == ASSISTANT_HEADER) + n\n",
    "\n",
    "\n",
    "def collate_fn(conversations):\n",
    "    batch = processor.apply_chat_template(conversations, tokenize=True, return_dict=True, return_tensors=\"pt\", processor_kwargs={\"padding\": True})\n",
    "    labels = torch.full_like(batch[\"input_ids\"], -100)\n",
    "    for i, ids in enumerate(batch[\"input_ids\"].tolist()):\n",
    "        start = answer_start(ids)\n",
    "        labels[i, start:] = batch[\"input_ids\"][i, start:]\n",
    "    labels[batch[\"attention_mask\"] == 0] = -100\n",
    "    batch[\"labels\"] = labels\n",
    "    return batch\n",
    "\n",
    "\n",
    "batch = collate_fn(train_dataset[:2])\n",
    "print({k: tuple(v.shape) for k, v in batch.items()})\n",
    "print(\"🎯 Supervised tokens:\", [processor.tokenizer.decode(ids[lab != -100]) for ids, lab in zip(batch[\"input_ids\"], batch[\"labels\"])])"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 📏 Baseline accuracy\n",
    "\n",
    "Let's measure how the model does before fine-tuning on the 200 held-out samples."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "import re\n",
    "\n",
    "\n",
    "@torch.inference_mode()\n",
    "def evaluate(model, show=3):\n",
    "    correct = 0\n",
    "    for k, sample in enumerate(tqdm(eval_raw, desc=\"Evaluating\")):\n",
    "        inputs = processor.apply_chat_template(\n",
    "            format_sample(sample, with_answer=False),\n",
    "            add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors=\"pt\",\n",
    "        ).to(model.device)\n",
    "        out = model.generate(**inputs, max_new_tokens=16, generation_mode=\"text\")\n",
    "        prediction = processor.decode(out.sequences[0, inputs[\"input_ids\"].shape[1]:], skip_special_tokens=True).strip()\n",
    "        match = re.search(r\"\\b([ABCD])\\b\", prediction)\n",
    "        correct += int(bool(match) and match.group(1) == sample[\"label\"])\n",
    "        if k < show:\n",
    "            print(f\"   {sample['question']!r} → {prediction!r} (answer {sample['label']})\")\n",
    "    return correct / len(eval_raw)\n",
    "\n",
    "\n",
    "accuracy_before = evaluate(model)\n",
    "print(f\"📏 Accuracy before fine-tuning: {accuracy_before:.1%}\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wrap the model with PEFT\n",
    "\n",
    "We add LoRA adapters to the language model: 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`). The vision and audio encoders stay frozen. LFM3's attention layers also have a per-head output gate, `gate_proj`; adding it to the targets is a reasonable choice. Feel free to modify the `r` (rank) value:\n",
    "- Higher → better approximation of full fine-tuning\n",
    "- Lower → needs even less compute resources"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from peft import LoraConfig, get_peft_model\n",
    "\n",
    "peft_config = LoraConfig(\n",
    "    r=16,\n",
    "    lora_alpha=32,\n",
    "    lora_dropout=0.05,\n",
    "    bias=\"none\",\n",
    "    task_type=\"CAUSAL_LM\",\n",
    "    # optional: add \"|gate_proj\" to the group for LFM3's per-head attention gate\n",
    "    target_modules=r\"model\\.language_model\\..*\\.(q_proj|k_proj|v_proj|o_proj|in_proj|out_proj|w1|w2|w3)\",\n",
    ")\n",
    "\n",
    "model = get_peft_model(model, peft_config)\n",
    "model.print_trainable_parameters()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🚀 Launch Training\n",
    "\n",
    "We use the 🤗 `Trainer`, which trains on the loss returned by the model. Our collate function builds the batches, so we keep all\n",
    "columns (`remove_unused_columns=False`). One epoch over the 1,800 training samples takes a few minutes on an A100."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from transformers import Trainer, TrainingArguments\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"lfm3-omni-med-vqa\",\n",
    "    num_train_epochs=1,\n",
    "    per_device_train_batch_size=2,  # images run at native resolution: keep this small on a 40 GB GPU\n",
    "    gradient_accumulation_steps=8,\n",
    "    learning_rate=2e-4,\n",
    "    lr_scheduler_type=\"cosine\",\n",
    "    warmup_steps=10,\n",
    "    weight_decay=0.01,\n",
    "    logging_steps=10,\n",
    "    save_strategy=\"no\",\n",
    "    bf16=True,\n",
    "    remove_unused_columns=False,\n",
    "    report_to=\"none\",\n",
    ")\n",
    "\n",
    "print(\"🏗️  Creating trainer...\")\n",
    "trainer = Trainer(model=model, args=training_args, train_dataset=train_dataset, data_collator=collate_fn)\n",
    "\n",
    "print(\"\\n🚀 Starting SFT training...\")\n",
    "trainer.train()\n",
    "print(\"🎉 SFT training completed!\")\n",
    "print(f\"💾 Peak GPU memory: {torch.cuda.max_memory_allocated() / 1e9:.1f} GB\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 📈 Accuracy after fine-tuning"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "model.eval()\n",
    "accuracy_after = evaluate(model)\n",
    "print(f\"📈 Accuracy  before: {accuracy_before:.1%}   after: {accuracy_after:.1%}\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 💾 Save the model\n",
    "\n",
    "Save the LoRA adapter, or merge it into the base weights to obtain a regular checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "model.save_pretrained(\"./lfm3-omni-med-vqa-lora\")\n",
    "print(\"💾 Adapter saved to: ./lfm3-omni-med-vqa-lora\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "print(\"🔄 Merging LoRA weights...\")\n",
    "merged = model.merge_and_unload()\n",
    "merged.save_pretrained(\"./lfm3-omni-med-vqa\")\n",
    "processor.save_pretrained(\"./lfm3-omni-med-vqa\")\n",
    "print(\"💾 Model saved to: ./lfm3-omni-med-vqa\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🚀 Next steps\n",
    "\n",
    "- **Real recordings**: synthetic questions are a quick way to bootstrap; mix in recordings from your users' microphones for robustness to accents and noise.\n",
    "- **Spoken answers**: to train the model to *answer* in speech, put the answer audio in the assistant turn; see [`lfm3_audio_tts_sft_with_transformers.ipynb`](lfm3_audio_tts_sft_with_transformers.ipynb) for how speech targets are built."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "A100",
   "provenance": [],
   "machine_shape": "hm"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}