{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cf52e279",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from peft import PeftModel\n",
    "from transformers import AutoProcessor, AutoModelForImageTextToText\n",
    "\n",
    "\n",
    "model_id = \"google/gemma-3-4b-pt\"\n",
    "output_dir = \"path-to-checkpoints-directory/google-gemma3-4b-pt/random-subset\"\n",
    "\n",
    "# Load Model base model\n",
    "model = AutoModelForImageTextToText.from_pretrained(model_id, low_cpu_mem_usage=True)\n",
    "\n",
    "# Merge LoRA and base model and save\n",
    "peft_model = PeftModel.from_pretrained(model, output_dir)\n",
    "merged_model = peft_model.merge_and_unload()\n",
    "merged_model.save_pretrained(os.path.join(output_dir, \"merged_model\"), safe_serialization=True, max_shard_size=\"2GB\")\n",
    "\n",
    "processor = AutoProcessor.from_pretrained(output_dir)\n",
    "processor.save_pretrained(os.path.join(output_dir, \"merged_model_processor\"))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}