{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9b09e756",
   "metadata": {},
   "source": [
    "# Teaching LFM3-VL-2B to Transcribe Formulas into LaTeX with TRL (GRPO)\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_vl_grpo_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 trains [LFM3-VL-2B](https://huggingface.co/LiquidAI/LFM3-VL-2B) to turn pictures of math formulas into LaTeX, with Group Relative Policy Optimization (GRPO) from [TRL](https://github.com/huggingface/trl) and LoRA. We measure the transcriptions before and after training on formulas the model never saw.\n",
    "\n",
    "GRPO needs no example answers to imitate, only a way to score the model's own answers. The model writes several answers to each picture, and it learns from the ones that score better than the others. LaTeX is a good fit: we can check the format of an answer and compare it with the reference automatically, and there are many ways to reason toward the right one.\n",
    "\n",
    "## 🎯 What you'll find\n",
    "\n",
    "- **The task**: a picture of a formula goes in, its LaTeX in a single `$$ ... $$` block comes out\n",
    "- **Two rewards**: one for the format of the answer, one for how close the LaTeX is to the reference\n",
    "- **Thinking kept on**: the model reasons before it answers, and GRPO also keeps its reasoning short\n",
    "- **Evaluation**: format rate, exact match, similarity, and output length, on held-out formulas"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddbd02aa",
   "metadata": {},
   "source": [
    "# 📦 Installation & Setup\n",
    "\n",
    "We pin every library to the versions this notebook was tested with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f32cd8ec",
   "metadata": {
    "tags": [
     "install"
    ]
   },
   "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": "8cb53a23",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\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-VL-2B 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": "1861000d",
   "metadata": {},
   "source": [
    "# Loading the model from Transformers 🤗\n",
    "\n",
    "LFM3-VL-2B 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": "6b4d2726",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForImageTextToText, AutoProcessor\n",
    "\n",
    "model_id = \"LiquidAI/LFM3-VL-2B\"\n",
    "\n",
    "print(\"📚 Loading processor...\")\n",
    "processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)\n",
    "processor.tokenizer.padding_side = \"left\"\n",
    "\n",
    "print(\"🧠 Loading model...\")\n",
    "model = AutoModelForImageTextToText.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": "7fe0c9a4",
   "metadata": {},
   "source": [
    "# 📚 Load the formula dataset\n",
    "\n",
    "[LaTeX_OCR](https://huggingface.co/datasets/unsloth/LaTeX_OCR) holds about 76,000 rendered formulas from physics papers, each with the LaTeX that produced it. The reference LaTeX puts a space between all tokens (`\\frac { 1 } { 2 }`) and uses the conventions of these papers: `{ \\cal R }` where recent LaTeX writes `\\mathcal { R }`, extra braces like `{ \\dot { \\phi } } ^ { 2 }`. The original model reads the formulas well but writes its own style, so GRPO has something to learn.\n",
    "\n",
    "GRPO learns from the model's own answers, so a few hundred pictures are enough: we train on 240 formulas and evaluate on 100 others."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b25da63",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"unsloth/LaTeX_OCR\").rename_column(\"text\", \"latex\")\n",
    "train_dataset = dataset[\"train\"].shuffle(seed=42).select(range(240))\n",
    "eval_dataset = dataset[\"test\"].shuffle(seed=42).select(range(100))\n",
    "print(dataset)\n",
    "\n",
    "example = train_dataset[0]\n",
    "print(f\"\\n🎯 Reference LaTeX:\\n{example['latex']}\")\n",
    "example[\"image\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c5624f8",
   "metadata": {},
   "source": [
    "## Build the prompts\n",
    "\n",
    "Each prompt is one user message with the picture and the instruction. GRPO only needs prompts: the reference LaTeX stays in its own column, and TRL passes it to the reward functions. TRL also puts the picture of the `image` column in front of the text.\n",
    "\n",
    "We keep LFM3-VL's reasoning on, as the chat template does by default: the model thinks, then answers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "306166d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "INSTRUCTION = \"Transcribe the formula in the picture into LaTeX. Answer with the LaTeX only, in a single $$ ... $$ block.\"\n",
    "\n",
    "def to_prompt(example):\n",
    "    return {\"prompt\": [{\"role\": \"user\", \"content\": INSTRUCTION}]}\n",
    "\n",
    "train_dataset = train_dataset.map(to_prompt)\n",
    "print(train_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1556f4b",
   "metadata": {},
   "source": [
    "# 🏆 Define the rewards\n",
    "\n",
    "TRL parses each completion into a message with the reasoning (`reasoning_content`) and the answer (`content`). A completion that runs out of tokens before the end of its reasoning has no answer, and gets no reward.\n",
    "\n",
    "- **Format** (0 or 1): the answer is exactly one `$$ ... $$` block, with nothing around it.\n",
    "- **LaTeX** (0 to 1): how close the LaTeX in the block is to the reference. We split both into LaTeX tokens (`\\frac`, `{`, `x`, ...), so spaces don't count, and compute `1 - edit distance / length`. An exact transcription scores 1.\n",
    "\n",
    "Both rewards also serve as evaluation metrics, with the exact match rate and the output length."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "175c1345",
   "metadata": {},
   "outputs": [],
   "source": [
    "BLOCK = re.compile(r\"\\s*\\$\\$([^$]+)\\$\\$\\s*\")  # <- the LaTeX inside holds no \"$\"\n",
    "TOKEN = re.compile(r\"\\\\[a-zA-Z]+|\\\\.|\\S\")\n",
    "\n",
    "def latex_of(answer):\n",
    "    \"\"\"The LaTeX in the answer's first $$ ... $$ block, or None.\"\"\"\n",
    "    match = BLOCK.search(answer or \"\")\n",
    "    return match.group(1) if match else None\n",
    "\n",
    "def similarity(prediction, reference):\n",
    "    \"\"\"1 - (edit distance between the LaTeX token sequences) / (length of the longer one).\"\"\"\n",
    "    a, b = TOKEN.findall(prediction), TOKEN.findall(reference)\n",
    "    previous = list(range(len(b) + 1))\n",
    "    for i, token_a in enumerate(a, 1):\n",
    "        current = [i]\n",
    "        for j, token_b in enumerate(b, 1):\n",
    "            current.append(min(previous[j] + 1, current[j - 1] + 1, previous[j - 1] + (token_a != token_b)))\n",
    "        previous = current\n",
    "    return 1 - previous[-1] / max(len(a), len(b), 1)\n",
    "\n",
    "def format_reward(completions, **kwargs):\n",
    "    return [float(BLOCK.fullmatch(completion[0].get(\"content\") or \"\") is not None) for completion in completions]\n",
    "\n",
    "def latex_reward(completions, latex, **kwargs):\n",
    "    rewards = []\n",
    "    for completion, reference in zip(completions, latex):\n",
    "        prediction = latex_of(completion[0].get(\"content\"))\n",
    "        rewards.append(similarity(prediction, reference) if prediction is not None else 0.0)\n",
    "    return rewards\n",
    "\n",
    "reference = \"\\\\frac { 1 } { 2 } m v ^ { 2 }\"\n",
    "for answer in [\"$$\\\\frac{1}{2} m v^{2}$$\", \"$$\\\\frac{1}{2} m v^2$$\", \"The formula is $$\\\\frac{1}{2}mv^{2}$$\", None]:\n",
    "    completion = [{\"role\": \"assistant\", \"content\": answer}]\n",
    "    print(f\"{answer!r:<42} format {format_reward([completion])[0]:.0f} | LaTeX {latex_reward([completion], [reference])[0]:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd2ba548",
   "metadata": {},
   "source": [
    "# 📏 Evaluate the model before training\n",
    "\n",
    "We transcribe the 100 evaluation formulas with the sampling settings recommended in the model card (they are the checkpoint's generation defaults), and report:\n",
    "\n",
    "- **Format**: the share of answers that are exactly one `$$ ... $$` block.\n",
    "- **Exact match**: the share of transcriptions with the same LaTeX tokens as the reference.\n",
    "- **Similarity**: the average LaTeX reward.\n",
    "- **Output tokens**: the average length of the reasoning and the answer, which sets the latency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "628c4c5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.auto import tqdm\n",
    "\n",
    "@torch.no_grad()\n",
    "def transcribe(model, images, max_new_tokens=2048):\n",
    "    \"\"\"The parsed message (reasoning and answer) of the model for each picture, and its length in tokens.\"\"\"\n",
    "    messages = [[{\"role\": \"user\", \"content\": [{\"type\": \"image\", \"image\": image}, {\"type\": \"text\", \"text\": INSTRUCTION}]}] for image in images]\n",
    "    inputs = processor.apply_chat_template(\n",
    "        messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors=\"pt\", processor_kwargs={\"padding\": True}\n",
    "    ).to(model.device)\n",
    "    outputs = model.generate(**inputs, max_new_tokens=max_new_tokens)[:, inputs[\"input_ids\"].shape[1]:]\n",
    "    results = []\n",
    "    for prompt_ids, output_ids in zip(inputs[\"input_ids\"], outputs):\n",
    "        output_ids = output_ids[output_ids != processor.tokenizer.pad_token_id].tolist()\n",
    "        prompt_ids = prompt_ids[prompt_ids != processor.tokenizer.pad_token_id].tolist()\n",
    "        results.append((processor.tokenizer.parse_response(output_ids, prefix=prompt_ids), len(output_ids)))\n",
    "    return results\n",
    "\n",
    "def evaluate(model, batch_size=50):\n",
    "    model.eval()\n",
    "    formats, exact, similarities, tokens = 0, 0, 0.0, 0\n",
    "    for start in tqdm(range(0, len(eval_dataset), batch_size), desc=\"Evaluating\"):\n",
    "        batch = eval_dataset.select(range(start, min(start + batch_size, len(eval_dataset))))\n",
    "        for (message, length), reference in zip(transcribe(model, batch[\"image\"]), batch[\"latex\"]):\n",
    "            completion = [message]\n",
    "            formats += format_reward([completion])[0]\n",
    "            reward = latex_reward([completion], [reference])[0]\n",
    "            similarities += reward\n",
    "            exact += reward == 1.0\n",
    "            tokens += length\n",
    "    n = len(eval_dataset)\n",
    "    return {\"format\": formats / n, \"exact_match\": exact / n, \"similarity\": similarities / n, \"tokens\": tokens / n}\n",
    "\n",
    "def show(r, title):\n",
    "    print(f\"{title}: format {r['format']:.1%} | exact match {r['exact_match']:.1%} | similarity {r['similarity']:.3f} | {r['tokens']:.0f} output tokens\")\n",
    "\n",
    "torch.manual_seed(0)\n",
    "baseline = evaluate(model)\n",
    "show(baseline, \"📏 Before training\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23dff9f9",
   "metadata": {},
   "source": [
    "# 🎯 GRPO training (LoRA)\n",
    "\n",
    "## Wrap the model with PEFT\n",
    "\n",
    "In LFM3's language model, 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. The vision encoder uses some of the same names, so we select the language model's layers with a regular expression and leave the vision encoder frozen."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b28b400",
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import LoraConfig\n",
    "\n",
    "peft_config = LoraConfig(\n",
    "    r=16,\n",
    "    lora_alpha=32,\n",
    "    # optional: add \"|gate_proj\" to the group for LFM3's per-head attention gate\n",
    "    target_modules=r\".*language_model.*\\.(q_proj|k_proj|v_proj|o_proj|in_proj|out_proj|w1|w2|w3)\",\n",
    "    task_type=\"CAUSAL_LM\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccb0a1f0",
   "metadata": {},
   "source": [
    "## Launch training\n",
    "\n",
    "At each step, the model writes 8 completions for each of 4 pictures. GRPO compares the rewards of the 8 completions of a picture: it makes the better ones more likely and the worse ones less likely.\n",
    "\n",
    "We sample with more randomness than the model card recommends for inference (`temperature=1.0` instead of 0.1): GRPO only learns when the completions of a picture differ. And we limit completions to 1,024 tokens: a completion that is still reasoning at the limit gets no reward, so the model learns to keep its reasoning short."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7062e73e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from trl import GRPOConfig, GRPOTrainer\n",
    "\n",
    "grpo_config = GRPOConfig(\n",
    "    output_dir=\"./LFM3-VL-2B-latex\",\n",
    "    max_steps=60,  # <- 60 steps × 4 pictures = one pass over the 240 training formulas\n",
    "    per_device_train_batch_size=8,\n",
    "    gradient_accumulation_steps=4,  # <- 32 completions per step: 4 pictures × 8 completions\n",
    "    num_generations=8,\n",
    "    max_completion_length=1024,\n",
    "    temperature=1.0,\n",
    "    learning_rate=5e-5,\n",
    "    lr_scheduler_type=\"constant_with_warmup\",\n",
    "    warmup_steps=0.05,  # <- a float is a fraction of the total steps\n",
    "    gradient_checkpointing=True,\n",
    "    logging_steps=5,\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 = GRPOTrainer(\n",
    "    model=model,\n",
    "    reward_funcs=[format_reward, latex_reward],\n",
    "    args=grpo_config,\n",
    "    train_dataset=train_dataset,\n",
    "    processing_class=processor,\n",
    "    peft_config=peft_config,\n",
    ")\n",
    "trainer.model.print_trainable_parameters()\n",
    "\n",
    "print(\"🚀 Starting GRPO training...\")\n",
    "trainer.train()\n",
    "print(\"🎉 GRPO training completed!\")\n",
    "print(f\"🧮 Peak GPU memory: {torch.cuda.max_memory_allocated() / 1e9:.1f} GB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b446dde4",
   "metadata": {},
   "source": [
    "## Training curves\n",
    "\n",
    "The rewards should go up and the completions should get shorter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99fbf119",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "logs = [entry for entry in trainer.state.log_history if \"reward\" in entry]\n",
    "steps = [entry[\"step\"] for entry in logs]\n",
    "figure, axes = plt.subplots(1, 3, figsize=(15, 3.5))\n",
    "for axis, key, title in zip(axes, [\"rewards/format_reward/mean\", \"rewards/latex_reward/mean\", \"completions/mean_length\"], [\"Format reward\", \"LaTeX reward\", \"Completion length (tokens)\"]):\n",
    "    axis.plot(steps, [entry[key] for entry in logs])\n",
    "    axis.set_title(title)\n",
    "    axis.set_xlabel(\"step\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b29298b",
   "metadata": {},
   "source": [
    "# 📊 Evaluate the trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "344f137f",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "trained = evaluate(trainer.model)\n",
    "show(baseline, \"📏 Before training\")\n",
    "show(trained, \"📏 After training\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5152ab71",
   "metadata": {},
   "source": [
    "# 🔍 Original vs trained model\n",
    "\n",
    "We transcribe a few evaluation formulas with the original model and the trained one, and render the three LaTeX strings. With LoRA, we don't need to load the original model again: `disable_adapter()` temporarily turns the adapter off."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da29b302",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Markdown, display\n",
    "\n",
    "trainer.model.eval()\n",
    "for example in eval_dataset.select(range(4)):\n",
    "    with trainer.model.disable_adapter():\n",
    "        [(original, original_length)] = transcribe(trainer.model, [example[\"image\"]])\n",
    "    [(trained_message, trained_length)] = transcribe(trainer.model, [example[\"image\"]])\n",
    "    display(example[\"image\"])\n",
    "    for name, message, length in [(\"Original\", original, original_length), (\"Trained\", trained_message, trained_length)]:\n",
    "        prediction = latex_of(message.get(\"content\"))\n",
    "        reward = latex_reward([[message]], [example[\"latex\"]])[0]\n",
    "        print(f\"{name} ({length} tokens, LaTeX reward {reward:.2f}): {prediction}\")\n",
    "    print(f\"Reference: {example['latex']}\")\n",
    "    display(Markdown(f\"Reference: $${example['latex']}$$\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba8e4ba1",
   "metadata": {},
   "source": [
    "And the summary over the evaluation formulas:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1152a77",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"{'':<16}{'original':>10}{'trained':>10}\")\n",
    "for metric in [\"format\", \"exact_match\"]:\n",
    "    print(f\"{metric:<16}{baseline[metric]:>10.1%}{trained[metric]:>10.1%}\")\n",
    "print(f\"{'similarity':<16}{baseline['similarity']:>10.3f}{trained['similarity']:>10.3f}\")\n",
    "print(f\"{'output tokens':<16}{baseline['tokens']:>10.0f}{trained['tokens']:>10.0f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "649e88f3",
   "metadata": {},
   "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-VL'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": "97117e6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_model(\"./LFM3-VL-2B-latex-lora\")\n",
    "\n",
    "output_dir = \"./LFM3-VL-2B-latex-merged\"\n",
    "merged_model = trainer.model.merge_and_unload()\n",
    "merged_model.save_pretrained(output_dir)\n",
    "processor.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-VL-2B-latex\")\n",
    "# processor.push_to_hub(\"your-username/LFM3-VL-2B-latex\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e43b3901",
   "metadata": {},
   "source": [
    "# 🚀 Next steps\n",
    "\n",
    "- **Your own checker**: any program that scores an answer is a reward: a LaTeX compiler, a unit test, a schema validator, a calculator.\n",
    "- **Longer training**: the dataset has 68,000 training formulas; more steps keep improving the rewards.\n",
    "- **Faster generation**: most of the training time goes into writing completions. With an inference engine that supports LFM3-VL, set `use_vllm=True` in `GRPOConfig`."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "A100",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}