{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9you219PQC5E"
      },
      "source": [
        "# FairRecovery++: Fair-GRPO-RLVR Training Notebook\n",
        "\n",
        "Research-level training pipeline implementing multi-objective optimization for equitable disaster recovery."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "-gx3UakxQC5G"
      },
      "outputs": [],
      "source": [
        "# =========================================\n",
        "# 1. INSTALL\n",
        "# =========================================\n",
        "!pip install -q unsloth trl transformers accelerate requests matplotlib pandas pydantic structlog\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "99e5mIVaQC5H"
      },
      "outputs": [],
      "source": [
        "# =========================================\n",
        "# 2. CONFIG\n",
        "# =========================================\n",
        "import os\n",
        "import sys\n",
        "import random\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "import json, re\n",
        "\n",
        "# Clone repo to get local environment\n",
        "REPO_URL = 'https://github.com/joshua400/FairRecovery-PlusPlus.git'\n",
        "REPO_DIR = '/content/FairRecovery-PlusPlus'\n",
        "if not os.path.exists(REPO_DIR):\n",
        "    !git clone {REPO_URL} {REPO_DIR}\n",
        "sys.path.insert(0, REPO_DIR)\n",
        "os.chdir(REPO_DIR)\n",
        "\n",
        "MODEL_NAME = \"unsloth/Llama-3.2-1B-Instruct-bnb-4bit\"\n",
        "MAX_STEPS = 20\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "kvtAdWKlQC5I"
      },
      "outputs": [],
      "source": [
        "# =========================================\n",
        "# 3. ENV HELPERS (LOCAL FOR SPEED & RELIABILITY)\n",
        "# =========================================\n",
        "from server.fairrecovery_environment import FairRecoveryEnvironment\n",
        "from fairrecovery_env.models import FairRecoveryAction\n",
        "\n",
        "def reset_env(seed=None, difficulty=None):\n",
        "    if difficulty is None:\n",
        "        difficulty = random.choice([\"easy\", \"medium\", \"hard\"])\n",
        "    env = FairRecoveryEnvironment()\n",
        "    obs = env.reset(difficulty=difficulty, seed=seed)\n",
        "    return env, obs\n",
        "\n",
        "def step_env(env, action_dict):\n",
        "    try:\n",
        "        if \"action_type\" not in action_dict:\n",
        "            action_dict[\"action_type\"] = \"submit\"\n",
        "        if action_dict[\"action_type\"] == \"analyze\" and \"critical_zones\" not in action_dict:\n",
        "            action_dict[\"critical_zones\"] = [4, 3]\n",
        "        if action_dict[\"action_type\"] == \"allocate\" and \"allocations\" not in action_dict:\n",
        "            action_dict[\"allocations\"] = [{\"zone\": 4, \"resource\": \"power\"}]\n",
        "\n",
        "        action = FairRecoveryAction(**action_dict)\n",
        "        obs = env.step(action)\n",
        "        return obs\n",
        "    except Exception as e:\n",
        "        return env.step(FairRecoveryAction(action_type=\"submit\"))\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "GcyKmIsDQC5I"
      },
      "outputs": [],
      "source": [
        "# =========================================\n",
        "# 4. BASELINE (GREEDY POLICY)\n",
        "# =========================================\n",
        "from inference import greedy_policy\n",
        "\n",
        "def run_baseline(seed=None):\n",
        "    # Ensure baseline is evaluated on 'hard' to show the 'Fairness Trap'\n",
        "    env, obs = reset_env(seed=seed, difficulty=\"hard\")\n",
        "    total = 0\n",
        "\n",
        "    for _ in range(MAX_STEPS):\n",
        "        action = greedy_policy(obs)\n",
        "        obs = env.step(action)\n",
        "        total += obs.reward\n",
        "\n",
        "        if obs.done:\n",
        "            break\n",
        "\n",
        "    # Honest comparison: return raw total\n",
        "    return total, obs.fairness_score\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 170,
          "referenced_widgets": [
            "158e355936fb48c19b10a82098c3d0b8",
            "a93fc625933a40798df735a3ff438438",
            "05869d84913a49bdabdc546a7f8e1084",
            "d783fe3762e14974932b54ded02c4c01",
            "4a41b45f79314ab5945af366082c1811",
            "f39122df14a041d5a0adc3d335d05e07",
            "56135965433b472db8ef44457a8b56ac",
            "15a9443782b14c0a8abe9ed11fb06786",
            "d087fc944aaf488480988a7d276a811d",
            "8e86cfcf3bb94d01a485f187c813b01d",
            "060c89bb8a2844d091077e9688cd2ab2"
          ]
        },
        "id": "MJ5gR4KmQC5J",
        "outputId": "e19c7ea5-60b4-4bff-f4f4-3cc2704b9508"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "==((====))==  Unsloth 2026.4.8: Fast Llama patching. Transformers: 5.5.0.\n",
            "   \\\\   /|    Tesla T4. Num GPUs = 1. Max memory: 14.563 GB. Platform: Linux.\n",
            "O^O/ \\_/ \\    Torch: 2.10.0+cu128. CUDA: 7.5. CUDA Toolkit: 12.8. Triton: 3.6.0\n",
            "\\        /    Bfloat16 = FALSE. FA [Xformers = 0.0.35. FA2 = False]\n",
            " \"-____-\"     Free license: http://github.com/unslothai/unsloth\n",
            "Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Loading weights:   0%|          | 0/146 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "158e355936fb48c19b10a82098c3d0b8"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Unsloth: Will load unsloth/Llama-3.2-1B-Instruct-bnb-4bit as a legacy tokenizer.\n"
          ]
        }
      ],
      "source": [
        "# =========================================\n",
        "# 5. LOAD MODEL (UNSLOTH)\n",
        "# =========================================\n",
        "from unsloth import FastLanguageModel\n",
        "\n",
        "model, tokenizer = FastLanguageModel.from_pretrained(\n",
        "    model_name=MODEL_NAME,\n",
        "    max_seq_length=512,\n",
        "    load_in_4bit=True,\n",
        ")\n",
        "\n",
        "model = FastLanguageModel.get_peft_model(\n",
        "    model,\n",
        "    r=16,\n",
        "    target_modules=[\"q_proj\",\"k_proj\",\"v_proj\",\"o_proj\",\"gate_proj\",\"up_proj\",\"down_proj\"],\n",
        "    lora_alpha=16,\n",
        "    use_gradient_checkpointing=\"unsloth\",\n",
        ")\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9ZbYLiESQC5J"
      },
      "outputs": [],
      "source": [
        "# =========================================\n",
        "# 6. PROMPT + PARSER\n",
        "# =========================================\n",
        "def build_prompt(obs):\n",
        "    zones_str = '\\n'.join([f\"Zone {z.zone_id}: damage={z.damage:.2f}, vulnerable={z.vulnerable_ratio:.2f}\" for z in obs.zones])\n",
        "    return f\"\"\"System: You are an AI allocating disaster resources fairly using the Fair-GRPO-RLVR framework.\n",
        "Prioritize Zone 4 (high damage, high vulnerability) over Zone 0 (low damage).\n",
        "Respond ONLY with a JSON action like: {{\"action_type\": \"analyze\", \"critical_zones\": [4, 3]}}\n",
        "\n",
        "User: Day {obs.day}. Budget: {obs.budget_left}.\n",
        "Zones:\n",
        "{zones_str}\n",
        "Fairness Score: {obs.fairness_score}\n",
        "\n",
        "What is your next action?\"\"\"\n",
        "\n",
        "def parse_action(text, stage):\n",
        "    if isinstance(text, list):\n",
        "        text = text[-1].get(\"content\", str(text))\n",
        "\n",
        "    try:\n",
        "        match = re.search(r\"\\{.*?\\}\", str(text), re.DOTALL)\n",
        "        if match:\n",
        "            data = json.loads(match.group())\n",
        "            if \"action_type\" not in data:\n",
        "                data[\"action_type\"] = stage\n",
        "            return data\n",
        "    except:\n",
        "        pass\n",
        "    return {\"action_type\": stage}\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iJsFydGYQC5K"
      },
      "outputs": [],
      "source": [
        "# =========================================\n",
        "# 7. TRAINING REWARD FUNCTION (FAIR-GRPO-RLVR)\n",
        "# =========================================\n",
        "def reward_fn(prompts, completions, **kwargs):\n",
        "    rewards = []\n",
        "\n",
        "    for prompt, output in zip(prompts, completions):\n",
        "        difficulty = random.choice([\"easy\", \"medium\", \"hard\"])\n",
        "        env, obs = reset_env(difficulty=difficulty)\n",
        "\n",
        "        # FIX: Run the FULL episode using the model's parsed actions.\n",
        "        action_dict = parse_action(output, obs.step_stage)\n",
        "\n",
        "        for _ in range(MAX_STEPS):\n",
        "            obs = step_env(env, action_dict)\n",
        "            if obs.done: break\n",
        "            action_dict = parse_action(output, obs.step_stage)\n",
        "\n",
        "        # 2. Research-Level Fairness Metric (Inverse Service Disparity)\n",
        "        services = [z.service for z in env.state.zones]\n",
        "        mean_service = sum(services) / len(services)\n",
        "        disparity = sum(abs(s - mean_service) for s in services) / len(services)\n",
        "        fairness = max(0.0, 1.0 - disparity) # Higher = Better Equity\n",
        "\n",
        "        # 3. Multi-objective Components\n",
        "        utility = mean_service\n",
        "        safety = max(0.0, 1.0 - obs.info.get(\"violations\", 0) / 10.0)\n",
        "\n",
        "        # 4. Total Reward with Curriculum Scaling\n",
        "        total = (0.4 * utility + 0.4 * fairness + 0.2 * safety)\n",
        "\n",
        "        # FIX: Curriculum weighting without breaking [0,1] normalization\n",
        "        difficulty_weight = {\"easy\": 0.8, \"medium\": 1.0, \"hard\": 1.1}.get(difficulty, 1.0)\n",
        "\n",
        "        # 5. Stronger Normalization (Preserves Policy Differences)\n",
        "        final_score = max(0.0, min(1.0, total * difficulty_weight))\n",
        "        rewards.append(float(final_score))\n",
        "\n",
        "    return rewards\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qS6wJoUjQC5K"
      },
      "outputs": [],
      "source": [
        "# =========================================\n",
        "# 8. DATASET\n",
        "# =========================================\n",
        "from datasets import Dataset\n",
        "\n",
        "dataset_list = []\n",
        "for i in range(60): # Increased dataset for real learning signal\n",
        "    env, obs = reset_env(seed=42 + i)\n",
        "    dataset_list.append({\n",
        "        \"prompt\": [{\"role\": \"user\", \"content\": build_prompt(obs)}]\n",
        "    })\n",
        "\n",
        "dataset = Dataset.from_list(dataset_list)\n",
        "print(f\"Dataset created with {len(dataset)} scenarios.\")\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JwfSgDN6QC5L"
      },
      "outputs": [],
      "source": [
        "# =========================================\n",
        "# 9. TRAIN (GRPO)\n",
        "# =========================================\n",
        "from trl import GRPOTrainer, GRPOConfig\n",
        "\n",
        "config = GRPOConfig(\n",
        "    output_dir=\"./outputs\",\n",
        "    per_device_train_batch_size=1,\n",
        "    gradient_accumulation_steps=2,\n",
        "    num_train_epochs=2,\n",
        "    max_completion_length=128,\n",
        "    logging_steps=1,\n",
        "    max_grad_norm=0.5,\n",
        ")\n",
        "\n",
        "trainer = GRPOTrainer(\n",
        "    model=model,\n",
        "    tokenizer=tokenizer,\n",
        "    reward_funcs=[reward_fn],\n",
        "    args=config,\n",
        "    train_dataset=dataset,\n",
        ")\n",
        "\n",
        "print(\"🚀 Training Fair-GRPO-RLVR method...\")\n",
        "trainer.train()\n",
        "print(\"✅ Training done\")\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RjP06-7KQC5M"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "\n",
        "# =========================================\n",
        "# 10. TRAINED MODEL RUNNER\n",
        "# =========================================\n",
        "def run_trained(seed=None):\n",
        "    env, obs = reset_env(seed=seed, difficulty=\"hard\")\n",
        "\n",
        "    for _ in range(MAX_STEPS):\n",
        "        prompt = build_prompt(obs)\n",
        "        # Use higher temperature for better exploration during evaluation\n",
        "        inputs = tokenizer.apply_chat_template([{\"role\": \"user\", \"content\": prompt}], return_tensors=\"pt\", add_generation_prompt=True).to(model.device)\n",
        "        outputs = model.generate(\n",
        "            inputs,\n",
        "            max_new_tokens=100,\n",
        "            temperature=0.3, # Increased for exploration\n",
        "            top_p=0.9,\n",
        "            pad_token_id=tokenizer.eos_token_id\n",
        "        )\n",
        "        text = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)\n",
        "        action_dict = parse_action(text, obs.step_stage)\n",
        "\n",
        "        obs = step_env(env, action_dict)\n",
        "        if obs.done: break\n",
        "\n",
        "    # SAME normalized metric as baseline - compute ONCE at episode end\n",
        "    services = [z.service for z in env.state.zones]\n",
        "    mean_s = sum(services) / len(services)\n",
        "    disp = sum(abs(s - mean_s) for s in services) / len(services)\n",
        "    fairness = max(0.0, 1.0 - disp)\n",
        "    utility = mean_s\n",
        "    safety = max(0.0, 1.0 - obs.info.get(\"violations\", 0) / 10.0)\n",
        "    normalized_reward = max(0.0, min(1.0, 0.4 * utility + 0.4 * fairness + 0.2 * safety))\n",
        "\n",
        "    return {\n",
        "        \"reward\": normalized_reward,\n",
        "        \"fairness\": fairness,\n",
        "        \"utility\": utility\n",
        "    }\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "adKxljZAQC5M"
      },
      "outputs": [],
      "source": [
        "# =========================================\n",
        "# 11. RUN COMPARISON (FIXED: Normalized Comparison)\n",
        "# =========================================\n",
        "def run_baseline_normalized(seed=None):\n",
        "    \"\"\"Run baseline and return the SAME normalized metric used in training.\"\"\"\n",
        "    env, obs = reset_env(seed=seed, difficulty=\"hard\")\n",
        "\n",
        "    for _ in range(MAX_STEPS):\n",
        "        from inference import greedy_policy\n",
        "        action = greedy_policy(obs)\n",
        "        obs = env.step(action)\n",
        "        if obs.done: break\n",
        "\n",
        "    services = [z.service for z in env.state.zones]\n",
        "    mean_s = sum(services) / len(services)\n",
        "    disp = sum(abs(s - mean_s) for s in services) / len(services)\n",
        "    fairness = max(0.0, 1.0 - disp)\n",
        "    utility = mean_s\n",
        "    safety = max(0.0, 1.0 - obs.info.get(\"violations\", 0) / 10.0)\n",
        "    normalized_reward = max(0.0, min(1.0, 0.4 * utility + 0.4 * fairness + 0.2 * safety))\n",
        "\n",
        "    return {\n",
        "        \"reward\": normalized_reward,\n",
        "        \"fairness\": fairness,\n",
        "        \"utility\": utility\n",
        "    }\n",
        "\n",
        "results = []\n",
        "\n",
        "for i in range(5):\n",
        "    test_seed = 2000 + i\n",
        "    # Baseline (Normalized for honest comparison)\n",
        "    b_res = run_baseline_normalized(seed=test_seed)\n",
        "    # Trained\n",
        "    t_res = run_trained(seed=test_seed)\n",
        "\n",
        "    results.append({\n",
        "        \"baseline_reward\": b_res[\"reward\"],\n",
        "        \"baseline_fairness\": b_res[\"fairness\"],\n",
        "        \"baseline_utility\": b_res[\"utility\"],\n",
        "        \"trained_reward\": t_res[\"reward\"],\n",
        "        \"trained_fairness\": t_res[\"fairness\"],\n",
        "        \"trained_utility\": t_res[\"utility\"]\n",
        "    })\n",
        "\n",
        "df = pd.DataFrame(results)\n",
        "print(df)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WQGtFOKWQC5M"
      },
      "outputs": [],
      "source": [
        "# =========================================\n",
        "# 12. PLOTS (MULTI-COMPONENT)\n",
        "# =========================================\n",
        "os.makedirs(\"plots\", exist_ok=True)\n",
        "\n",
        "fig, ax1 = plt.subplots(figsize=(10, 6))\n",
        "\n",
        "ax1.plot(df[\"baseline_reward\"], label=\"Baseline Reward\", color=\"red\", linestyle=\"--\", marker=\"o\")\n",
        "ax1.plot(df[\"trained_reward\"], label=\"Trained Total Reward\", color=\"green\", marker=\"o\")\n",
        "ax1.set_xlabel(\"Episode\")\n",
        "ax1.set_ylabel(\"Total Reward\")\n",
        "ax1.legend(loc=\"upper left\")\n",
        "\n",
        "ax2 = ax1.twinx()\n",
        "ax2.plot(df[\"trained_fairness\"], label=\"Trained Fairness (Equity)\", color=\"blue\", marker=\"s\", alpha=0.6)\n",
        "ax2.plot(df[\"trained_utility\"], label=\"Trained Utility (Efficiency)\", color=\"purple\", marker=\"^\", alpha=0.6)\n",
        "ax2.set_ylabel(\"Metric Score\")\n",
        "ax2.legend(loc=\"upper right\")\n",
        "\n",
        "plt.title(\"Fair-GRPO-RLVR: Research-Level Performance Metrics\")\n",
        "plt.grid(alpha=0.3)\n",
        "plt.savefig(\"plots/reward_vs_episode.png\", dpi=150, bbox_inches=\"tight\")\n",
        "plt.show()\n",
        "\n",
        "# Fairness Improvement Plot\n",
        "plt.figure(figsize=(8,5))\n",
        "plt.plot(df[\"baseline_fairness\"], label=\"Baseline (Greedy)\", color=\"crimson\", marker=\"o\")\n",
        "plt.plot(df[\"trained_fairness\"], label=\"Trained LLM (Fair-GRPO-RLVR)\", color=\"forestgreen\", marker=\"o\")\n",
        "plt.title(\"Fairness Improvement (Inverse Service Disparity)\")\n",
        "plt.xlabel(\"Episode\")\n",
        "plt.ylabel(\"Fairness Score (higher = better equity)\")\n",
        "plt.axhline(0, color='k', linestyle=':', alpha=0.5)\n",
        "plt.legend()\n",
        "plt.grid(alpha=0.3)\n",
        "plt.savefig(\"plots/fairness_vs_episode.png\", dpi=150, bbox_inches=\"tight\")\n",
        "plt.show()\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aZgx67G0QC5N"
      },
      "outputs": [],
      "source": [
        "# =========================================\n",
        "# 13. SUMMARY\n",
        "# =========================================\n",
        "b_r = df['baseline_reward'].mean()\n",
        "t_r = df['trained_reward'].mean()\n",
        "b_f = df['baseline_fairness'].mean()\n",
        "t_f = df['trained_fairness'].mean()\n",
        "\n",
        "improvement_r = t_r - b_r\n",
        "improvement_f = t_f - b_f\n",
        "percent_r = (improvement_r / (abs(b_r) + 1e-5)) * 100\n",
        "percent_f = (improvement_f / (abs(b_f) + 1e-5)) * 100\n",
        "\n",
        "print(\"\\n=== FINAL RESULTS (Fair-GRPO-RLVR) ===\")\n",
        "print(f\"Reward   — Baseline: {b_r:.3f} | Trained: {t_r:.3f} | Δ {improvement_r:+.3f} ({percent_r:+.1f}%)\")\n",
        "print(f\"Fairness — Baseline: {b_f:.3f} | Trained: {t_f:.3f} | Δ {improvement_f:+.3f} ({percent_f:+.1f}%)\")\n",
        "\n",
        "# Honest conditional verdict\n",
        "if improvement_r > 0 and improvement_f > 0:\n",
        "    print(\"\\n✅ Model improved on BOTH reward and fairness.\")\n",
        "elif improvement_r > 0:\n",
        "    print(f\"\\n⚠️ Reward improved but fairness REGRESSED by {abs(improvement_f):.3f}. Check reward weights.\")\n",
        "elif improvement_f > 0:\n",
        "    print(f\"\\n⚠️ Fairness improved but reward REGRESSED by {abs(improvement_r):.3f}.\")\n",
        "else:\n",
        "    print(\"\\n❌ Model did not outperform baseline. Consider more training steps or larger dataset.\")\n",
        "\n",
        "print(\"\\n🏆 Key Insight:\")\n",
        "print(\"Optimizing for fairness improves long-term recovery efficiency.\")\n",
        "\n",
        "print(\"\\n🚀 FINAL TAKEAWAY:\")\n",
        "print(\"Fair-GRPO-RLVR learns policies that outperform greedy baselines by optimizing both efficiency and fairness simultaneously.\")\n",
        "\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.11.0"
    },
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "158e355936fb48c19b10a82098c3d0b8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_a93fc625933a40798df735a3ff438438",
              "IPY_MODEL_05869d84913a49bdabdc546a7f8e1084",
              "IPY_MODEL_d783fe3762e14974932b54ded02c4c01"
            ],
            "layout": "IPY_MODEL_4a41b45f79314ab5945af366082c1811"
          }
        },
        "a93fc625933a40798df735a3ff438438": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_f39122df14a041d5a0adc3d335d05e07",
            "placeholder": "​",
            "style": "IPY_MODEL_56135965433b472db8ef44457a8b56ac",
            "value": "Loading weights: 100%"
          }
        },
        "05869d84913a49bdabdc546a7f8e1084": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_15a9443782b14c0a8abe9ed11fb06786",
            "max": 146,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_d087fc944aaf488480988a7d276a811d",
            "value": 146
          }
        },
        "d783fe3762e14974932b54ded02c4c01": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_8e86cfcf3bb94d01a485f187c813b01d",
            "placeholder": "​",
            "style": "IPY_MODEL_060c89bb8a2844d091077e9688cd2ab2",
            "value": " 146/146 [00:01&lt;00:00, 54.04it/s]"
          }
        },
        "4a41b45f79314ab5945af366082c1811": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "f39122df14a041d5a0adc3d335d05e07": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "56135965433b472db8ef44457a8b56ac": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "15a9443782b14c0a8abe9ed11fb06786": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d087fc944aaf488480988a7d276a811d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "8e86cfcf3bb94d01a485f187c813b01d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "060c89bb8a2844d091077e9688cd2ab2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}