Finetune_llava-v1.6-mistral-7b-SYDNEYFISH_A100.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LLaVA Fine-tuning on Sydney Fish Dataset (A100 Optimized)\n",
    "\n",
    "This notebook implements:\n",
    "1. Stratified sampling across fish species\n",
    "2. A100-optimized training parameters\n",
    "3. Comprehensive testing and evaluation\n",
    "4. Visualization of correct and incorrect predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import BitsAndBytesConfig, LlavaNextForConditionalGeneration, AutoProcessor\n",
    "import lightning as L\n",
    "from torch.utils.data import DataLoader\n",
    "import re\n",
    "from nltk import edit_distance\n",
    "import numpy as np\n",
    "from lightning.pytorch.callbacks import Callback\n",
    "from lightning.pytorch.callbacks.early_stopping import EarlyStopping\n",
    "from huggingface_hub import notebook_login, HfApi\n",
    "from Sydney_Fish_Dataset_Stratified import StratifiedSydneyFishDataset\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from sklearn.metrics import classification_report, confusion_matrix\n",
    "import json\n",
    "from IPython.display import display\n",
    "from PIL import Image\n",
    "\n",
    "# Enable tensor cores for better performance\n",
    "torch.set_float32_matmul_precision('high')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LLaVA Fine-tuning on Sydney Fish Dataset (A100 Optimized)\n",
    "\n",
    "This notebook implements:\n",
    "1. Stratified sampling across fish species\n",
    "2. A100-optimized training parameters\n",
    "3. Comprehensive testing and evaluation\n",
    "4. Visualization of correct and incorrect predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import BitsAndBytesConfig, LlavaNextForConditionalGeneration, AutoProcessor\n",
    "import lightning as L\n",
    "from torch.utils.data import DataLoader\n",
    "import re\n",
    "from nltk import edit_distance\n",
    "import numpy as np\n",
    "from lightning.pytorch.callbacks import Callback\n",
    "from lightning.pytorch.callbacks.early_stopping import EarlyStopping\n",
    "from huggingface_hub import notebook_login, HfApi\n",
    "from Sydney_Fish_Dataset_Stratified import StratifiedSydneyFishDataset\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from sklearn.metrics import classification_report, confusion_matrix\n",
    "import json\n",
    "from IPython.display import display\n",
    "from PIL import Image\n",
    "\n",
    "# Enable tensor cores for better performance\n",
    "torch.set_float32_matmul_precision('high')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Login to Hugging Face\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check GPU and memory\n",
    "if torch.cuda.is_available():\n",
    "    print(\"PyTorch is connected to GPU.\")\n",
    "    print(f\"GPU Device Name: {torch.cuda.get_device_name(0)}\")\n",
    "    print(f\"Number of GPUs available: {torch.cuda.device_count()}\")\n",
    "    print(f\"Current GPU: {torch.cuda.current_device()}\")\n",
    "    print(f\"Available GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB\")\n",
    "else:\n",
    "    print(\"PyTorch is not connected to GPU.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configuration\n",
    "MAX_LENGTH = 256  # Can be larger on A100\n",
    "MODEL_ID = \"llava-hf/llava-v1.6-mistral-7b-hf\"\n",
    "REPO_ID = \"YOUR_HUGGINGFACE_USERNAME/llava-v1.6-mistral-7b-sydneyfish-a100\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load processor and model\n",
    "processor = AutoProcessor.from_pretrained(MODEL_ID)\n",
    "processor.tokenizer.padding_side = \"right\"\n",
    "\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=torch.float16,\n",
    "    bnb_4bit_use_double_quant=True\n",
    ")\n",
    "\n",
    "model = LlavaNextForConditionalGeneration.from_pretrained(\n",
    "    MODEL_ID,\n",
    "    torch_dtype=torch.float16,\n",
    "    quantization_config=bnb_config,\n",
    "    device_map=\"auto\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply PEFT with A100-optimized settings\n",
    "from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model\n",
    "\n",
    "def find_all_linear_names(model):\n",
    "    cls = torch.nn.Linear\n",
    "    lora_module_names = set()\n",
    "    multimodal_keywords = ['multi_modal_projector', 'vision_model']\n",
    "    \n",
    "    for name, module in model.named_modules():\n",
    "        if any(mm_keyword in name for mm_keyword in multimodal_keywords):\n",
    "            continue\n",
    "        if isinstance(module, cls):\n",
    "            names = name.split('.')\n",
    "            lora_module_names.add(names[0] if len(names) == 1 else names[-1])\n",
    "    \n",
    "    if 'lm_head' in lora_module_names:\n",
    "        lora_module_names.remove('lm_head')\n",
    "    return list(lora_module_names)\n",
    "\n",
    "lora_config = LoraConfig(\n",
    "    r=16,  # Increased for A100\n",
    "    lora_alpha=32,  # Increased for A100\n",
    "    lora_dropout=0.1,\n",
    "    target_modules=find_all_linear_names(model),\n",
    "    init_lora_weights=\"gaussian\",\n",
    ")\n",
    "\n",
    "model = prepare_model_for_kbit_training(model)\n",
    "model = get_peft_model(model, lora_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load datasets with stratified sampling\n",
    "train_dataset = StratifiedSydneyFishDataset(split=\"train\", seed=42)\n",
    "val_dataset = StratifiedSydneyFishDataset(split=\"validation\", seed=42)\n",
    "test_dataset = StratifiedSydneyFishDataset(split=\"test\", seed=42)\n",
    "\n",
    "print(\"\\nDataset Statistics:\")\n",
    "print(f\"Training samples: {len(train_dataset)}\")\n",
    "print(f\"Validation samples: {len(val_dataset)}\")\n",
    "print(f\"Test samples: {len(test_dataset)}\")\n",
    "\n",
    "print(\"\\nTraining Set Species Distribution:\")\n",
    "for species, count in sorted(train_dataset.get_species_distribution().items()):\n",
    "    print(f\"{species}: {count} samples\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LlavaModelPLModule(L.LightningModule):\n",
    "    def __init__(self, config, processor, model):\n",
    "        super().__init__()\n",
    "        self.config = config\n",
    "        self.processor = processor\n",
    "        self.model = model\n",
    "        self.batch_size = config.get(\"batch_size\")\n",
    "        self.test_predictions = []\n",
    "        self.test_targets = []\n",
    "        self.test_images = []\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        input_ids, attention_mask, pixel_values, image_sizes, labels = batch\n",
    "        outputs = self.model(\n",
    "            input_ids=input_ids,\n",
    "            attention_mask=attention_mask,\n",
    "            pixel_values=pixel_values,\n",
    "            image_sizes=image_sizes,\n",
    "            labels=labels\n",
    "        )\n",
    "        loss = outputs.loss\n",
    "        self.log(\"train_loss\", loss)\n",
    "        return loss\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        input_ids, attention_mask, pixel_values, image_sizes, labels = batch\n",
    "        outputs = self.model(\n",
    "            input_ids=input_ids,\n",
    "            attention_mask=attention_mask,\n",
    "            pixel_values=pixel_values,\n",
    "            image_sizes=image_sizes,\n",
    "            labels=labels\n",
    "        )\n",
    "        loss = outputs.loss\n",
    "        self.log(\"val_loss\", loss)\n",
    "        return loss\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        input_ids, attention_mask, pixel_values, image_sizes, labels = batch\n",
    "        \n",
    "        # Store original images for visualization\n",
    "        self.test_images.extend(pixel_values.cpu().numpy())\n",
    "        \n",
    "        # Generate predictions\n",
    "        generated_ids = self.model.generate(\n",
    "            input_ids=input_ids,\n",
    "            attention_mask=attention_mask,\n",
    "            pixel_values=pixel_values,\n",
    "            image_sizes=image_sizes,\n",
    "            max_new_tokens=MAX_LENGTH\n",
    "        )\n",
    "        \n",
    "        predictions = self.processor.batch_decode(generated_ids[:, input_ids.size(1):], skip_special_tokens=True)\n",
    "        targets = self.processor.batch_decode(labels, skip_special_tokens=True)\n",
    "        \n",
    "        self.test_predictions.extend(predictions)\n",
    "        self.test_targets.extend(targets)\n",
    "        \n",
    "        return {\"predictions\": predictions, \"targets\": targets}\n",
    "\n",
    "    def on_test_epoch_end(self):\n",
    "        # Extract species names\n",
    "        pred_species = []\n",
    "        target_species = []\n",
    "        correct_indices = []\n",
    "        incorrect_indices = []\n",
    "        \n",
    "        for i, (pred, target) in enumerate(zip(self.test_predictions, self.test_targets)):\n",
    "            try:\n",
    "                pred_json = json.loads(pred)\n",
    "                pred_name = pred_json[\"species\"][\"name\"]\n",
    "                pred_species.append(pred_name)\n",
    "            except:\n",
    "                pred_species.append(\"ERROR\")\n",
    "                \n",
    "            try:\n",
    "                target_json = json.loads(target)\n",
    "                target_name = target_json[\"species\"][\"name\"]\n",
    "                target_species.append(target_name)\n",
    "            except:\n",
    "                target_species.append(\"ERROR\")\n",
    "            \n",
    "            if pred_species[-1] == target_species[-1]:\n",
    "                correct_indices.append(i)\n",
    "            else:\n",
    "                incorrect_indices.append(i)\n",
    "        \n",
    "        # Calculate and display metrics\n",
    "        accuracy = len(correct_indices) / len(self.test_predictions)\n",
    "        print(f\"\\nTest Accuracy: {accuracy * 100:.2f}%\")\n",
    "        \n",
    "        print(\"\\nClassification Report:\")\n",
    "        print(classification_report(target_species, pred_species))\n",
    "        \n",
    "        # Plot confusion matrix\n",
    "        plt.figure(figsize=(15, 10))\n",
    "        cm = confusion_matrix(target_species, pred_species)\n",
    "        sns.heatmap(cm, annot=True, fmt='d', \n",
    "                    xticklabels=sorted(set(target_species)), \n",
    "                    yticklabels=sorted(set(target_species)))\n",
    "        plt.title('Confusion Matrix')\n",
    "        plt.xlabel('Predicted')\n",
    "        plt.ylabel('True')\n",
    "        plt.xticks(rotation=45)\n",
    "        plt.yticks(rotation=45)\n",
    "        plt.tight_layout()\n",
    "        plt.show()\n",
    "        \n",
    "        # Display example predictions\n",
    "        print(\"\\nExample Correct Predictions:\")\n",
    "        for idx in correct_indices[:5]:\n",
    "            print(f\"\\nTrue: {target_species[idx]}\")\n",
    "            print(f\"Predicted: {pred_species[idx]}\")\n",
    "            img = self.test_images[idx]\n",
    "            plt.imshow(np.transpose(img, (1, 2, 0)))\n",
    "            plt.axis('off')\n",
    "            plt.show()\n",
    "            \n",
    "        print(\"\\nExample Incorrect Predictions:\")\n",
    "        for idx in incorrect_indices[:5]:\n",
    "            print(f\"\\nTrue: {target_species[idx]}\")\n",
    "            print(f\"Predicted: {pred_species[idx]}\")\n",
    "            img = self.test_images[idx]\n",
    "            plt.imshow(np.transpose(img, (1, 2, 0)))\n",
    "            plt.axis('off')\n",
    "            plt.show()\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.AdamW(self.parameters(), lr=self.config.get(\"lr\"))\n",
    "        return optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A100-optimized training configuration\n",
    "config = {\n",
    "    \"max_epochs\": 10,\n",
    "    \"check_val_every_n_epoch\": 1,\n",
    "    \"gradient_clip_val\": 1.0,\n",
    "    \"accumulate_grad_batches\": 4,  # Reduced for A100\n",
    "    \"lr\": 2e-4,\n",
    "    \"batch_size\": 8,  # Increased for A100\n",
    "    \"num_nodes\": 1,\n",
    "    \"warmup_steps\": 100,\n",
    "    \"result_path\": \"./result\",\n",
    "    \"verbose\": True,\n",
    "    \"num_workers\": 8  # Increased for A100\n",
    "}\n",
    "\n",
    "model_module = LlavaModelPLModule(config, processor, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Callbacks\n",
    "class PushToHubCallback(Callback):\n",
    "    def on_train_epoch_end(self, trainer, pl_module):\n",
    "        print(f\"Pushing model to the hub, epoch {trainer.current_epoch}\")\n",
    "        pl_module.model.push_to_hub(\n",
    "            REPO_ID,\n",
    "            commit_message=f\"Training in progress, epoch {trainer.current_epoch}\"\n",
    "        )\n",
    "\n",
    "    def on_train_end(self, trainer, pl_module):\n",
    "        print(f\"Pushing model to the hub after training\")\n",
    "        pl_module.processor.push_to_hub(\n",
    "            REPO_ID,\n",
    "            commit_message=f\"Training done\"\n",
    "        )\n",
    "        pl_module.model.push_to_hub(\n",
    "            REPO_ID,\n",
    "            commit_message=f\"Training done\"\n",
    "        )\n",
    "\n",
    "early_stop_callback = EarlyStopping(\n",
    "    monitor=\"val_loss\",\n",
    "    patience=3,\n",
    "    verbose=True,\n",
    "    mode=\"min\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize trainer\n",
    "trainer = L.Trainer(\n",
    "    accelerator=\"gpu\",\n",
    "    devices=[0],\n",
    "    max_epochs=config.get(\"max_epochs\"),\n",
    "    accumulate_grad_batches=config.get(\"accumulate_grad_batches\"),\n",
    "    check_val_every_n_epoch=config.get(\"check_val_every_n_epoch\"),\n",
    "    gradient_clip_val=config.get(\"gradient_clip_val\"),\n",
    "    precision=\"16-mixed\",\n",
    "    callbacks=[\n",
    "        PushToHubCallback(),\n",
    "        early_stop_callback\n",
    "    ],\n",
    "    strategy=\"auto\",\n",
    "    enable_progress_bar=True,\n",
    "    enable_model_summary=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the model\n",
    "trainer.fit(model_module)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test the model and visualize results\n",
    "trainer.test(model_module)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Login to Hugging Face\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check GPU and memory\n",
    "if torch.cuda.is_available():\n",
    "    print(\"PyTorch is connected to GPU.\")\n",
    "    print(f\"GPU Device Name: {torch.cuda.get_device_name(0)}\")\n",
    "    print(f\"Number of GPUs available: {torch.cuda.device_count()}\")\n",
    "    print(f\"Current GPU: {torch.cuda.current_device()}\")\n",
    "    print(f\"Available GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB\")\n",
    "else:\n",
    "    print(\"PyTorch is not connected to GPU.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configuration\n",
    "MAX_LENGTH = 256  # Can be larger on A100\n",
    "MODEL_ID = \"llava-hf/llava-v1.6-mistral-7b-hf\"\n",
    "REPO_ID = \"YOUR_HUGGINGFACE_USERNAME/llava-v1.6-mistral-7b-sydneyfish-a100\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load processor and model\n",
    "processor = AutoProcessor.from_pretrained(MODEL_ID)\n",
    "processor.tokenizer.padding_side = \"right\"\n",
    "\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=torch.float16,\n",
    "    bnb_4bit_use_double_quant=True\n",
    ")\n",
    "\n",
    "model = LlavaNextForConditionalGeneration.from_pretrained(\n",
    "    MODEL_ID,\n",
    "    torch_dtype=torch.float16,\n",
    "    quantization_config=bnb_config,\n",
    "    device_map=\"auto\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply PEFT with A100-optimized settings\n",
    "from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model\n",
    "\n",
    "def find_all_linear_names(model):\n",
    "    cls = torch.nn.Linear\n",
    "    lora_module_names = set()\n",
    "    multimodal_keywords = ['multi_modal_projector', 'vision_model']\n",
    "    \n",
    "    for name, module in model.named_modules():\n",
    "        if any(mm_keyword in name for mm_keyword in multimodal_keywords):\n",
    "            continue\n",
    "        if isinstance(module, cls):\n",
    "            names = name.split('.')\n",
    "            lora_module_names.add(names[0] if len(names) == 1 else names[-1])\n",
    "    \n",
    "    if 'lm_head' in lora_module_names:\n",
    "        lora_module_names.remove('lm_head')\n",
    "    return list(lora_module_names)\n",
    "\n",
    "lora_config = LoraConfig(\n",
    "    r=16,  # Increased for A100\n",
    "    lora_alpha=32,  # Increased for A100\n",
    "    lora_dropout=0.1,\n",
    "    target_modules=find_all_linear_names(model),\n",
    "    init_lora_weights=\"gaussian\",\n",
    ")\n",
    "\n",
    "model = prepare_model_for_kbit_training(model)\n",
    "model = get_peft_model(model, lora_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load datasets with stratified sampling\n",
    "train_dataset = StratifiedSydneyFishDataset(split=\"train\", seed=42)\n",
    "val_dataset = StratifiedSydneyFishDataset(split=\"validation\", seed=42)\n",
    "test_dataset = StratifiedSydneyFishDataset(split=\"test\", seed=42)\n",
    "\n",
    "print(\"\\nDataset Statistics:\")\n",
    "print(f\"Training samples: {len(train_dataset)}\")\n",
    "print(f\"Validation samples: {len(val_dataset)}\")\n",
    "print(f\"Test samples: {len(test_dataset)}\")\n",
    "\n",
    "print(\"\\nTraining Set Species Distribution:\")\n",
    "for species, count in sorted(train_dataset.get_species_distribution().items()):\n",
    "    print(f\"{species}: {count} samples\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LlavaModelPLModule(L.LightningModule):\n",
    "    def __init__(self, config, processor, model):\n",
    "        super().__init__()\n",
    "        self.config = config\n",
    "        self.processor = processor\n",
    "        self.model = model\n",
    "        self.batch_size = config.get(\"batch_size\")\n",
    "        self.test_predictions = []\n",
    "        self.test_targets = []\n",
    "        self.test_images = []\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        input_ids, attention_mask, pixel_values, image_sizes, labels = batch\n",
    "        outputs = self.model(\n",
    "            input_ids=input_ids,\n",
    "            attention_mask=attention_mask,\n",
    "            pixel_values=pixel_values,\n",
    "            image_sizes=image_sizes,\n",
    "            labels=labels\n",
    "        )\n",
    "        loss = outputs.loss\n",
    "        self.log(\"train_loss\", loss)\n",
    "        return loss\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        input_ids, attention_mask, pixel_values, image_sizes, labels = batch\n",
    "        outputs = self.model(\n",
    "            input_ids=input_ids,\n",
    "            attention_mask=attention_mask,\n",
    "            pixel_values=pixel_values,\n",
    "            image_sizes=image_sizes,\n",
    "            labels=labels\n",
    "        )\n",
    "        loss = outputs.loss\n",
    "        self.log(\"val_loss\", loss)\n",
    "        return loss\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        input_ids, attention_mask, pixel_values, image_sizes, labels = batch\n",
    "        \n",
    "        # Store original images for visualization\n",
    "        self.test_images.extend(pixel_values.cpu().numpy())\n",
    "        \n",
    "        # Generate predictions\n",
    "        generated_ids = self.model.generate(\n",
    "            input_ids=input_ids,\n",
    "            attention_mask=attention_mask,\n",
    "            pixel_values=pixel_values,\n",
    "            image_sizes=image_sizes,\n",
    "            max_new_tokens=MAX_LENGTH\n",
    "        )\n",
    "        \n",
    "        predictions = self.processor.batch_decode(generated_ids[:, input_ids.size(1):], skip_special_tokens=True)\n",
    "        targets = self.processor.batch_decode(labels, skip_special_tokens=True)\n",
    "        \n",
    "        self.test_predictions.extend(predictions)\n",
    "        self.test_targets.extend(targets)\n",
    "        \n",
    "        return {\"predictions\": predictions, \"targets\": targets}\n",
    "\n",
    "    def on_test_epoch_end(self):\n",
    "        # Extract species names\n",
    "        pred_species = []\n",
    "        target_species = []\n",
    "        correct_indices = []\n",
    "        incorrect_indices = []\n",
    "        \n",
    "        for i, (pred, target) in enumerate(zip(self.test_predictions, self.test_targets)):\n",
    "            try:\n",
    "                pred_json = json.loads(pred)\n",
    "                pred_name = pred_json[\"species\"][\"name\"]\n",
    "                pred_species.append(pred_name)\n",
    "            except:\n",
    "                pred_species.append(\"ERROR\")\n",
    "                \n",
    "            try:\n",
    "                target_json = json.loads(target)\n",
    "                target_name = target_json[\"species\"][\"name\"]\n",
    "                target_species.append(target_name)\n",
    "            except:\n",
    "                target_species.append(\"ERROR\")\n",
    "            \n",
    "            if pred_species[-1] == target_species[-1]:\n",
    "                correct_indices.append(i)\n",
    "            else:\n",
    "                incorrect_indices.append(i)\n",
    "        \n",
    "        # Calculate and display metrics\n",
    "        accuracy = len(correct_indices) / len(self.test_predictions)\n",
    "        print(f\"\\nTest Accuracy: {accuracy * 100:.2f}%\")\n",
    "        \n",
    "        print(\"\\nClassification Report:\")\n",
    "        print(classification_report(target_species, pred_species))\n",
    "        \n",
    "        # Plot confusion matrix\n",
    "        plt.figure(figsize=(15, 10))\n",
    "        cm = confusion_matrix(target_species, pred_species)\n",
    "        sns.heatmap(cm, annot=True, fmt='d', \n",
    "                    xticklabels=sorted(set(target_species)), \n",
    "                    yticklabels=sorted(set(target_species)))\n",
    "        plt.title('Confusion Matrix')\n",
    "        plt.xlabel('Predicted')\n",
    "        plt.ylabel('True')\n",
    "        plt.xticks(rotation=45)\n",
    "        plt.yticks(rotation=45)\n",
    "        plt.tight_layout()\n",
    "        plt.show()\n",
    "        \n",
    "        # Display example predictions\n",
    "        print(\"\\nExample Correct Predictions:\")\n",
    "        for idx in correct_indices[:5]:\n",
    "            print(f\"\\nTrue: {target_species[idx]}\")\n",
    "            print(f\"Predicted: {pred_species[idx]}\")\n",
    "            img = self.test_images[idx]\n",
    "            plt.imshow(np.transpose(img, (1, 2, 0)))\n",
    "            plt.axis('off')\n",
    "            plt.show()\n",
    "            \n",
    "        print(\"\\nExample Incorrect Predictions:\")\n",
    "        for idx in incorrect_indices[:5]:\n",
    "            print(f\"\\nTrue: {target_species[idx]}\")\n",
    "            print(f\"Predicted: {pred_species[idx]}\")\n",
    "            img = self.test_images[idx]\n",
    "            plt.imshow(np.transpose(img, (1, 2, 0)))\n",
    "            plt.axis('off')\n",
    "            plt.show()\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.AdamW(self.parameters(), lr=self.config.get(\"lr\"))\n",
    "        return optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A100-optimized training configuration\n",
    "config = {\n",
    "    \"max_epochs\": 10,\n",
    "    \"check_val_every_n_epoch\": 1,\n",
    "    \"gradient_clip_val\": 1.0,\n",
    "    \"accumulate_grad_batches\": 4,  # Reduced for A100\n",
    "    \"lr\": 2e-4,\n",
    "    \"batch_size\": 8,  # Increased for A100\n",
    "    \"num_nodes\": 1,\n",
    "    \"warmup_steps\": 100,\n",
    "    \"result_path\": \"./result\",\n",
    "    \"verbose\": True,\n",
    "    \"num_workers\": 8  # Increased for A100\n",
    "}\n",
    "\n",
    "model_module = LlavaModelPLModule(config, processor, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Callbacks\n",
    "class PushToHubCallback(Callback):\n",
    "    def on_train_epoch_end(self, trainer, pl_module):\n",
    "        print(f\"Pushing model to the hub, epoch {trainer.current_epoch}\")\n",
    "        pl_module.model.push_to_hub(\n",
    "            REPO_ID,\n",
    "            commit_message=f\"Training in progress, epoch {trainer.current_epoch}\"\n",
    "        )\n",
    "\n",
    "    def on_train_end(self, trainer, pl_module):\n",
    "        print(f\"Pushing model to the hub after training\")\n",
    "        pl_module.processor.push_to_hub(\n",
    "            REPO_ID,\n",
    "            commit_message=f\"Training done\"\n",
    "        )\n",
    "        pl_module.model.push_to_hub(\n",
    "            REPO_ID,\n",
    "            commit_message=f\"Training done\"\n",
    "        )\n",
    "\n",
    "early_stop_callback = EarlyStopping(\n",
    "    monitor=\"val_loss\",\n",
    "    patience=3,\n",
    "    verbose=True,\n",
    "    mode=\"min\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize trainer\n",
    "trainer = L.Trainer(\n",
    "    accelerator=\"gpu\",\n",
    "    devices=[0],\n",
    "    max_epochs=config.get(\"max_epochs\"),\n",
    "    accumulate_grad_batches=config.get(\"accumulate_grad_batches\"),\n",
    "    check_val_every_n_epoch=config.get(\"check_val_every_n_epoch\"),\n",
    "    gradient_clip_val=config.get(\"gradient_clip_val\"),\n",
    "    precision=\"16-mixed\",\n",
    "    callbacks=[\n",
    "        PushToHubCallback(),\n",
    "        early_stop_callback\n",
    "    ],\n",
    "    strategy=\"auto\",\n",
    "    enable_progress_bar=True,\n",
    "    enable_model_summary=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the model\n",
    "trainer.fit(model_module)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test the model and visualize results\n",
    "trainer.test(model_module)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}