diff --git a/Diabetic_Retinopathy_Classification.ipynb b/Diabetic_Retinopathy_Classification.ipynb
new file mode 100644
index 0000000..e878325
--- /dev/null
+++ b/Diabetic_Retinopathy_Classification.ipynb
@@ -0,0 +1,295 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Diabetic Retinopathy Classification\n",
+    "\n",
+    "This notebook fine-tunes a model for 'Referral' / 'No Referral' diabetic retinopathy classification using the APTOS 2019 and MESSIDOR-2 datasets."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Environment Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "# --- Install Dependencies ---\n",
+    "print(\"\\n⏳ Installing dependencies...\")\n",
+    "!pip install torch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 --index-url https://download.pytorch.org/whl/cu121 -q\n",
+    "!pip install timm==0.9.16 pandas==2.2.2 scikit-learn -q\n",
+    "!pip install gdown -q\n",
+    "print(\"✅ Dependencies installed.\")\n",
+    "\n",
+    "# --- Set up device ---\n",
+    "import torch\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "print(f\"Using device: {device}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Download and Preprocess Datasets"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import gdown\n",
+    "import os\n",
+    "\n",
+    "# --- Download APTOS 2019 ---\n",
+    "aptos_url = 'https://drive.google.com/uc?id=162YPf4OhMVxj9TrQH0GnJv0n7z7gJWpj'\n",
+    "aptos_output = 'APTOS2019.zip'\n",
+    "if not os.path.exists(aptos_output):\n",
+    "    print('Downloading APTOS 2019 dataset...')\n",
+    "    gdown.download(aptos_url, aptos_output, quiet=False)\n",
+    "    !unzip -q {aptos_output}\n",
+    "    print('APTOS 2019 dataset downloaded and unzipped.')\n",
+    "else:\n",
+    "    print('APTOS 2019 dataset already downloaded.')\n",
+    "\n",
+    "# --- Download MESSIDOR-2 ---\n",
+    "messidor_url = 'https://drive.google.com/uc?id=1vOLBUK9xdzNV8eVkRjVdNrRwhPfaOmda'\n",
+    "messidor_output = 'MESSIDOR2.zip'\n",
+    "if not os.path.exists(messidor_output):\n",
+    "    print('Downloading MESSIDOR-2 dataset...')\n",
+    "    gdown.download(messidor_url, messidor_output, quiet=False)\n",
+    "    !unzip -q {messidor_output}\n",
+    "    print('MESSIDOR-2 dataset downloaded and unzipped.')\n",
+    "else:\n",
+    "    print('MESSIDOR-2 dataset already downloaded.')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "from sklearn.model_selection import train_test_split\n",
+    "\n",
+    "# --- Preprocess APTOS 2019 ---\n",
+    "aptos_df = pd.read_csv('train.csv')\n",
+    "aptos_df['image_path'] = aptos_df['id_code'].apply(lambda x: os.path.join('train_images', x + '.png'))\n",
+    "\n",
+    "# --- Preprocess MESSIDOR-2 ---\n",
+    "messidor_df = pd.read_csv('messidor_data.csv')\n",
+    "messidor_df['image_path'] = messidor_df['image_id'].apply(lambda x: os.path.join('messidor-2', 'images', x + '.jpg'))\n",
+    "\n",
+    "# --- Combine datasets ---\n",
+    "combined_df = pd.concat([\n",
+    "    aptos_df[['image_path', 'diagnosis']],\n",
+    "    messidor_df[['image_path', 'adjudicated_dr_grade']]\n",
+    "], ignore_index=True)\n",
+    "combined_df.rename(columns={'diagnosis': 'grade', 'adjudicated_dr_grade': 'grade'}, inplace=True)\n",
+    "\n",
+    "# --- Create binary labels ---\n",
+    "# Referral: grade >= 2\n",
+    "# No Referral: grade < 2\n",
+    "combined_df['label'] = combined_df['grade'].apply(lambda x: 1 if x >= 2 else 0)\n",
+    "\n",
+    "# --- Split data ---\n",
+    "train_df, val_df = train_test_split(combined_df, test_size=0.2, stratify=combined_df['label'], random_state=42)\n",
+    "\n",
+    "print(f'Training samples: {len(train_df)}')\n",
+    "print(f'Validation samples: {len(val_df)}')\n",
+    "print(train_df['label'].value_counts())\n",
+    "print(val_df['label'].value_counts())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Implement Data Loading"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from torchvision import transforms\n",
+    "from PIL import Image\n",
+    "\n",
+    "class DRDataset(Dataset):\n",
+    "    def __init__(self, df, transform=None):\n",
+    "        self.df = df\n",
+    "        self.transform = transform\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.df)\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "        image_path = self.df.iloc[idx]['image_path']\n",
+    "        image = Image.open(image_path).convert('RGB')\n",
+    "        label = self.df.iloc[idx]['label']\n",
+    "\n",
+    "        if self.transform:\n",
+    "            image = self.transform(image)\n",
+    "\n",
+    "        return image, label\n",
+    "\n",
+    "data_transforms = {\n",
+    "    'train': transforms.Compose([\n",
+    "        transforms.Resize((256, 256)),\n",
+    "        transforms.RandomHorizontalFlip(),\n",
+    "        transforms.RandomRotation(10),\n",
+    "        transforms.ToTensor(),\n",
+    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
+    "    ]),\n",
+    "    'val': transforms.Compose([\n",
+    "        transforms.Resize((256, 256)),\n",
+    "        transforms.ToTensor(),\n",
+    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
+    "    ]),\n",
+    "}\n",
+    "\n",
+    "train_dataset = DRDataset(train_df, transform=data_transforms['train'])\n",
+    "val_dataset = DRDataset(val_df, transform=data_transforms['val'])\n",
+    "\n",
+    "train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=2)\n",
+    "val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=2)\n",
+    "\n",
+    "print(f'Train loader: {len(train_loader)} batches')\n",
+    "print(f'Validation loader: {len(val_loader)} batches')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Enhance Fine-Tuning Process"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from torchvision import models\n",
+    "\n",
+    "# --- Define the model ---\n",
+    "model = models.resnet50(pretrained=True)\n",
+    "num_ftrs = model.fc.in_features\n",
+    "model.fc = nn.Linear(num_ftrs, 2) # Binary classification\n",
+    "model = model.to(device)\n",
+    "\n",
+    "# --- Define loss function and optimizer ---\n",
+    "criterion = nn.CrossEntropyLoss()\n",
+    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
+    "\n",
+    "print(\"Model, loss function, and optimizer are ready.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Train and Evaluate the Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score\n",
+    "import numpy as np\n",
+    "\n",
+    "def train_one_epoch(model, loader, criterion, optimizer, device):\n",
+    "    model.train()\n",
+    "    running_loss = 0.0\n",
+    "    for images, labels in loader:\n",
+    "        images, labels = images.to(device), labels.to(device)\n",
+    "\n",
+    "        optimizer.zero_grad()\n",
+    "        outputs = model(images)\n",
+    "        loss = criterion(outputs, labels)\n",
+    "        loss.backward()\n",
+    "        optimizer.step()\n",
+    "\n",
+    "        running_loss += loss.item() * images.size(0)\n",
+    "\n",
+    "    epoch_loss = running_loss / len(loader.dataset)\n",
+    "    return epoch_loss\n",
+    "\n",
+    "def evaluate_model(model, loader, criterion, device):\n",
+    "    model.eval()\n",
+    "    running_loss = 0.0\n",
+    "    all_preds, all_labels = [], []\n",
+    "\n",
+    "    with torch.no_grad():\n",
+    "        for images, labels in loader:\n",
+    "            images, labels = images.to(device), labels.to(device)\n",
+    "\n",
+    "            outputs = model(images)\n",
+    "            loss = criterion(outputs, labels)\n",
+    "            running_loss += loss.item() * images.size(0)\n",
+    "\n",
+    "            _, preds = torch.max(outputs, 1)\n",
+    "            all_preds.extend(preds.cpu().numpy())\n",
+    "            all_labels.extend(labels.cpu().numpy())\n",
+    "\n",
+    "    epoch_loss = running_loss / len(loader.dataset)\n",
+    "    accuracy = accuracy_score(all_labels, all_preds)\n",
+    "    precision = precision_score(all_labels, all_preds)\n",
+    "    recall = recall_score(all_labels, all_preds)\n",
+    "    f1 = f1_score(all_labels, all_preds)\n",
+    "\n",
+    "    return epoch_loss, accuracy, precision, recall, f1\n",
+    "\n",
+    "num_epochs = 10\n",
+    "best_accuracy = 0.0\n",
+    "\n",
+    "for epoch in range(num_epochs):\n",
+    "    train_loss = train_one_epoch(model, train_loader, criterion, optimizer, device)\n",
+    "    val_loss, val_acc, val_prec, val_rec, val_f1 = evaluate_model(model, val_loader, criterion, device)\n",
+    "\n",
+    "    print(f'Epoch {epoch+1}/{num_epochs} | Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f} | Val Precision: {val_prec:.4f} | Val Recall: {val_rec:.4f} | Val F1: {val_f1:.4f}')\n",
+    "\n",
+    "    if val_acc > best_accuracy:\n",
+    "        best_accuracy = val_acc\n",
+    "        torch.save(model.state_dict(), 'best_model.pth')\n",
+    "        print('Best model saved.')"
+   ]
+  }
+ ],
+ "metadata": {
+  "colab": {
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}
diff --git a/RETFound_Reproduction.ipynb b/RETFound_Reproduction.ipynb
new file mode 100644
index 0000000..09b1146
--- /dev/null
+++ b/RETFound_Reproduction.ipynb
@@ -0,0 +1,451 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# RETFound: Fine-Tuning and Evaluation for Diabetic Retinopathy\n",
+    "\n",
+    "This notebook provides a complete, end-to-end workflow to reproduce the fine-tuning of the RETFound model on the IDRiD dataset for diabetic retinopathy classification."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# ==============================================================================\n",
+    "# ## 1. Environment Setup\n",
+    "#\n",
+    "# This section clones the original RETFound_MAE repository, installs all\n",
+    "# necessary dependencies, and applies essential patches to the scripts to ensure\n",
+    "# they run correctly in the Colab environment.\n",
+    "# =============================================================================="
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "# --- Clone Repository & Set Working Directory ---\n",
+    "repo_dir = 'RETFound_MAE'\n",
+    "\n",
+    "if not os.path.isdir(repo_dir):\n",
+    "  print(\"Cloning repository...\")\n",
+    "  !git clone https://github.com/rmaphoh/RETFound_MAE/\n",
+    "else:\n",
+    "  print(\"Repository already exists. Skipping clone.\")\n",
+    "\n",
+    "%cd {repo_dir}\n",
+    "\n",
+    "# --- Install Dependencies ---\n",
+    "print(\"\\n⏳ Installing dependencies...\")\n",
+    "!pip install torch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 --index-url https://download.pytorch.org/whl/cu121 -q\n",
+    "!pip install timm==0.9.16 pandas==2.2.2 scikit-learn -q\n",
+    "print(\"✅ Dependencies installed.\")\n",
+    "\n",
+    "# --- Apply Initial Patches ---\n",
+    "# Note: Further patches are applied programmatically in the next section.\n",
+    "print(\"\\n⚙️ Patching script to bypass errors...\")\n",
+    "!sed -i \"s/if True:  # args.distributed:/if args.distributed:/g\" main_finetune.py\n",
+    "print(\"✅ Script patched.\")\n",
+    "\n",
+    "print(\"\\n🎉 Setup complete. Ready for the next step.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# ==============================================================================\n",
+    "# ## 2. Generate Core Scripts\n",
+    "#\n",
+    "# We programmatically overwrite `main_finetune.py` and `engine_finetune.py`\n",
+    "# to incorporate all the fixes and improvements we developed. This includes\n",
+    "# robust metric calculation (AUROC, Specificity), error handling, and disabling\n",
+    "# conflicting libraries.\n",
+    "# =============================================================================="
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# --- Create main_finetune.py ---\n",
+    "main_finetune_content = r\"\"\"\n",
+    "import argparse, datetime, json, numpy as np, os, time\n",
+    "from pathlib import Path\n",
+    "import torch, torch.backends.cudnn as cudnn\n",
+    "from timm.data.mixup import Mixup\n",
+    "import models_vit as models, util.lr_decay as lrd, util.misc as misc\n",
+    "from util.datasets import build_dataset\n",
+    "from util.misc import NativeScalerWithGradNormCount as NativeScaler\n",
+    "from huggingface_hub import hf_hub_download\n",
+    "from engine_finetune import train_one_epoch, evaluate\n",
+    "import warnings\n",
+    "\n",
+    "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
+    "\n",
+    "def get_args_parser():\n",
+    "    parser = argparse.ArgumentParser('MAE fine-tuning for image classification', add_help=False)\n",
+    "    parser.add_argument('--batch_size', default=128, type=int)\n",
+    "    parser.add_argument('--epochs', default=50, type=int)\n",
+    "    parser.add_argument('--accum_iter', default=1, type=int)\n",
+    "    parser.add_argument('--model', default='vit_large_patch16', type=str)\n",
+    "    parser.add_argument('--input_size', default=256, type=int)\n",
+    "    parser.add_argument('--drop_path', type=float, default=0.2)\n",
+    "    parser.add_argument('--clip_grad', type=float, default=None)\n",
+    "    parser.add_argument('--weight_decay', type=float, default=0.05)\n",
+    "    parser.add_argument('--lr', type=float, default=None)\n",
+    "    parser.add_argument('--blr', type=float, default=5e-3)\n",
+    "    parser.add_argument('--layer_decay', type=float, default=0.65)\n",
+    "    parser.add_argument('--min_lr', type=float, default=1e-6)\n",
+    "    parser.add_argument('--warmup_epochs', type=int, default=10)\n",
+    "    parser.add_argument('--finetune', default='', type=str)\n",
+    "    parser.add_argument('--task', default='', type=str)\n",
+    "    parser.add_argument('--global_pool', action='store_true', default=True)\n",
+    "    parser.add_argument('--data_path', default='./data/', type=str)\n",
+    "    parser.add_argument('--nb_classes', default=5, type=int)\n",
+    "    parser.add_argument('--output_dir', default='./output_dir')\n",
+    "    parser.add_argument('--log_dir', default='./output_logs')\n",
+    "    parser.add_argument('--device', default='cuda')\n",
+    "    parser.add_argument('--seed', default=0, type=int)\n",
+    "    parser.add_argument('--resume', default='')\n",
+    "    parser.add_argument('--start_epoch', default=0, type=int)\n",
+    "    parser.add_argument('--eval', action='store_true')\n",
+    "    parser.add_argument('--num_workers', default=2, type=int)\n",
+    "    parser.add_argument('--pin_mem', action='store_true', default=True)\n",
+    "    parser.add_argument('--world_size', default=1, type=int)\n",
+    "    parser.add_argument('--local_rank', default=-1, type=int)\n",
+    "    parser.add_argument('--dist_on_itp', action='store_true')\n",
+    "    parser.add_argument('--dist_url', default='env://')\n",
+    "    return parser\n",
+    "\n",
+    "def main(args):\n",
+    "    misc.init_distributed_mode(args)\n",
+    "    device = torch.device(args.device)\n",
+    "    seed = args.seed + misc.get_rank()\n",
+    "    torch.manual_seed(seed)\n",
+    "    np.random.seed(seed)\n",
+    "    cudnn.benchmark = True\n",
+    "\n",
+    "    dataset_test = build_dataset(is_train='test', args=args)\n",
+    "    if not args.eval:\n",
+    "        dataset_train = build_dataset(is_train='train', args=args)\n",
+    "        dataset_val = build_dataset(is_train='val', args=args)\n",
+    "    else:\n",
+    "        dataset_train, dataset_val = None, None\n",
+    "\n",
+    "    if args.distributed:\n",
+    "        # ... (distributed setup omitted for Colab clarity)\n",
+    "        pass\n",
+    "    else:\n",
+    "        sampler_test = torch.utils.data.SequentialSampler(dataset_test)\n",
+    "        if not args.eval:\n",
+    "            sampler_train = torch.utils.data.RandomSampler(dataset_train)\n",
+    "            sampler_val = torch.utils.data.SequentialSampler(dataset_val)\n",
+    "\n",
+    "    log_writer = None # Disabled for Colab\n",
+    "\n",
+    "    data_loader_test = torch.utils.data.DataLoader(dataset_test, sampler=sampler_test, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False)\n",
+    "    if not args.eval:\n",
+    "        data_loader_train = torch.utils.data.DataLoader(dataset_train, sampler=sampler_train, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True)\n",
+    "        data_loader_val = torch.utils.data.DataLoader(dataset_val, sampler=sampler_val, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False)\n",
+    "\n",
+    "    model = models.__dict__[args.model](img_size=args.input_size, num_classes=args.nb_classes, drop_path_rate=args.drop_path, global_pool=args.global_pool)\n",
+    "\n",
+    "    if args.finetune and not args.eval:\n",
+    "        print(f\"Downloading pre-trained weights from Hugging Face: {args.finetune}\")\n",
+    "        checkpoint_path = hf_hub_download(repo_id=f'YukunZhou/{args.finetune}', filename=f'{args.finetune}.pth')\n",
+    "        checkpoint = torch.load(checkpoint_path, map_location='cpu')['model']\n",
+    "        msg = model.load_state_dict(checkpoint, strict=False)\n",
+    "        print(f\"Loaded pre-trained checkpoint from {args.finetune} with message: {msg}\")\n",
+    "\n",
+    "    if args.resume:\n",
+    "        print(f\"Resuming from checkpoint: {args.resume}\")\n",
+    "        checkpoint = torch.load(args.resume, map_location='cpu')['model']\n",
+    "        model.load_state_dict(checkpoint, strict=False)\n",
+    "\n",
+    "    model.to(device)\n",
+    "    print(f'Number of model params (M): {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1.e6:.2f}')\n",
+    "\n",
+    "    if args.eval:\n",
+    "        evaluate(data_loader_test, model, device, args, 0, 'test', args.nb_classes, log_writer)\n",
+    "        return\n",
+    "\n",
+    "    eff_batch_size = args.batch_size * misc.get_world_size()\n",
+    "    if args.lr is None: args.lr = args.blr * eff_batch_size / 256\n",
+    "    print(f\"Actual lr: {args.lr:.2e}\")\n",
+    "    param_groups = lrd.param_groups_lrd(model, args.weight_decay, no_weight_decay_list=model.no_weight_decay(), layer_decay=args.layer_decay)\n",
+    "    optimizer = torch.optim.AdamW(param_groups, lr=args.lr)\n",
+    "    loss_scaler = NativeScaler()\n",
+    "    criterion = torch.nn.CrossEntropyLoss()\n",
+    "\n",
+    "    print(f\"--- Starting Training for {args.epochs} epochs ---\")\n",
+    "    max_accuracy = 0.0\n",
+    "    for epoch in range(args.start_epoch, args.epochs):\n",
+    "        train_stats = train_one_epoch(model, criterion, data_loader_train, optimizer, device, epoch, loss_scaler, args=args)\n",
+    "        val_stats, _ = evaluate(data_loader_val, model, device, args, epoch, 'val', args.nb_classes, log_writer=log_writer)\n",
+    "        print(f\"EPOCH:{epoch} | Val Acc: {val_stats['acc1']:.1f}%\")\n",
+    "        if max_accuracy < val_stats[\"acc1\"]:\n",
+    "            max_accuracy = val_stats[\"acc1\"]\n",
+    "            misc.save_model(args=args, model=model, model_without_ddp=model, optimizer=optimizer, loss_scaler=loss_scaler, epoch=epoch, mode='best')\n",
+    "        print(f'Max accuracy: {max_accuracy:.2f}%')\n",
+    "\n",
+    "if __name__ == '__main__':\n",
+    "    args = get_args_parser().parse_args()\n",
+    "    if args.output_dir: Path(os.path.join(args.output_dir, args.task)).mkdir(parents=True, exist_ok=True)\n",
+    "    main(args)\n",
+    "\"\"\"\n",
+    "with open(\"main_finetune.py\", \"w\") as f:\n",
+    "    f.write(main_finetune_content)\n",
+    "print(\"✅ `main_finetune.py` generated.\")\n",
+    "\n",
+    "\n",
+    "# --- Create engine_finetune.py ---\n",
+    "engine_finetune_content = r\"\"\"\n",
+    "import math, sys\n",
+    "from typing import Iterable, Optional\n",
+    "import torch, torch.nn.functional as F\n",
+    "import numpy as np\n",
+    "from sklearn.metrics import roc_auc_score, classification_report, confusion_matrix\n",
+    "from timm.utils import accuracy\n",
+    "import util.misc as misc, util.lr_sched as lr_sched\n",
+    "\n",
+    "def train_one_epoch(model: torch.nn.Module, criterion: torch.nn.Module,\n",
+    "                    data_loader: Iterable, optimizer: torch.optim.Optimizer,\n",
+    "                    device: torch.device, epoch: int, loss_scaler, args=None):\n",
+    "    model.train(True)\n",
+    "    metric_logger = misc.MetricLogger(delimiter=\"  \")\n",
+    "    metric_logger.add_meter('lr', misc.SmoothedValue(window_size=1, fmt='{value:.6f}'))\n",
+    "    header = f'Epoch: [{epoch}]'\n",
+    "    for data_iter_step, (samples, targets) in enumerate(metric_logger.log_every(data_loader, 20, header)):\n",
+    "        lr_sched.adjust_learning_rate(optimizer, data_iter_step / len(data_loader) + epoch, args)\n",
+    "        samples, targets = samples.to(device, non_blocking=True), targets.to(device, non_blocking=True)\n",
+    "        with torch.cuda.amp.autocast():\n",
+    "            loss = criterion(model(samples), targets)\n",
+    "        loss_value = loss.item()\n",
+    "        if not math.isfinite(loss_value): sys.exit(f\"Loss is {loss_value}, stopping training\")\n",
+    "        loss_scaler(loss, optimizer, parameters=model.parameters())\n",
+    "        optimizer.zero_grad()\n",
+    "        torch.cuda.synchronize()\n",
+    "        metric_logger.update(loss=loss_value)\n",
+    "        metric_logger.update(lr=optimizer.param_groups[0][\"lr\"])\n",
+    "    print(\"Averaged stats:\", metric_logger)\n",
+    "    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}\n",
+    "\n",
+    "@torch.no_grad()\n",
+    "def evaluate(data_loader, model, device, args, epoch, mode, num_class, log_writer=None):\n",
+    "    criterion = torch.nn.CrossEntropyLoss()\n",
+    "    metric_logger = misc.MetricLogger(delimiter=\"  \")\n",
+    "    header, model.eval(), all_preds, all_labels, all_probs = 'Test:', [], [], []\n",
+    "    for batch in metric_logger.log_every(data_loader, 10, header):\n",
+    "        images, target = batch[0].to(device, non_blocking=True), batch[-1].to(device, non_blocking=True)\n",
+    "        with torch.cuda.amp.autocast():\n",
+    "            output = model(images)\n",
+    "            loss = criterion(output, target)\n",
+    "        preds, probs = torch.argmax(output, dim=1), F.softmax(output, dim=1)\n",
+    "        all_preds.extend(preds.cpu().numpy())\n",
+    "        all_labels.extend(target.cpu().numpy())\n",
+    "        all_probs.extend(probs.cpu().detach().numpy())\n",
+    "        metric_logger.update(loss=loss.item())\n",
+    "        metric_logger.meters['acc1'].update(accuracy(output, target, topk=(1,))[0].item(), n=images.shape[0])\n",
+    "    print(f'* Acc@1 {metric_logger.acc1.global_avg:.3f} loss {metric_logger.loss.global_avg:.3f}')\n",
+    "    all_labels, all_preds, all_probs = np.array(all_labels), np.array(all_preds), np.array(all_probs)\n",
+    "    print(\"\\n--- Performance Metrics ---\")\n",
+    "    if len(np.unique(all_labels)) > 1 and len(np.unique(all_preds)) > 1:\n",
+    "        try:\n",
+    "            print(f\"AUROC (Macro): {roc_auc_score(all_labels, all_probs, multi_class='ovr', average='macro'):.4f}\")\n",
+    "        except Exception as e: print(f\"Could not calculate AUROC: {e}\")\n",
+    "    else: print(\"Skipping AUROC: not enough classes in labels or predictions.\")\n",
+    "    cm, fp = confusion_matrix(all_labels, all_preds), cm.sum(axis=0) - np.diag(cm)\n",
+    "    tn = cm.sum() - (fp + (cm.sum(axis=1) - np.diag(cm)) + np.diag(cm))\n",
+    "    print(f\"Specificity (Macro): {np.mean(tn / (tn + fp)):.4f}\")\n",
+    "    print(\"\\n--- Classification Report ---\")\n",
+    "    print(classification_report(all_labels, all_preds, target_names=[f'Class {i}' for i in range(num_class)], digits=4, zero_division=0))\n",
+    "    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}, 'results'\n",
+    "\"\"\"\n",
+    "with open(\"engine_finetune.py\", \"w\") as f:\n",
+    "    f.write(engine_finetune_content)\n",
+    "print(\"✅ `engine_finetune.py` generated.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# ==============================================================================\n",
+    "# ## 3. Data Preparation\n",
+    "#\n",
+    "# Download and unzip the IDRiD dataset for 5-class DR grading.\n",
+    "# =============================================================================="
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data_folder_path = '../IDRiD_data'\n",
+    "\n",
+    "if not os.path.isdir(data_folder_path):\n",
+    "    print(\"⏳ Dataset folder not found. Downloading and unzipping...\")\n",
+    "    !gdown --id 1c6zexA705z-ANEBNXJOBsk6uCvRnzmr3 -O ../IDRiD_data.zip\n",
+    "    !unzip -q -o ../IDRiD_data.zip -d ../\n",
+    "    print(\"✅ Dataset downloaded and ready.\")\n",
+    "else:\n",
+    "    print(f\"✅ Dataset already exists at '{data_folder_path}'. Skipping download.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# ==============================================================================\n",
+    "# ## 4. Workflow A: Fine-Tune and Evaluate\n",
+    "#\n",
+    "# This is the primary workflow. It involves authenticating with Hugging Face,\n",
+    "# running the fine-tuning process on the IDRiD dataset, and then evaluating\n",
+    "# the best model produced during that run.\n",
+    "# =============================================================================="
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# --- Step 4.1: Authenticate with Hugging Face ---\n",
+    "from huggingface_hub import notebook_login\n",
+    "print(\"🚀 Please log in to Hugging Face to download the base model.\")\n",
+    "notebook_login()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# --- Step 4.2: Run Fine-Tuning ---\n",
+    "import torch\n",
+    "from pathlib import Path\n",
+    "from main_finetune import get_args_parser, main as finetune_main\n",
+    "\n",
+    "print(\"\\n🚀 Starting fine-tuning...\")\n",
+    "ft_parser = get_args_parser()\n",
+    "args_ft = ft_parser.parse_args([\n",
+    "    '--model', 'RETFound_mae',\n",
+    "    '--epochs', '20',\n",
+    "    '--blr', '5e-3',\n",
+    "    '--data_path', '../IDRiD_data',\n",
+    "    '--task', 'RETFound_finetune_IDRiD',\n",
+    "    '--output_dir', './output_dir',\n",
+    "    '--log_dir', './log_dir',\n",
+    "    '--finetune', 'RETFound_mae_meh'\n",
+    "])\n",
+    "Path(os.path.join(args_ft.output_dir, args_ft.task)).mkdir(parents=True, exist_ok=True)\n",
+    "finetune_main(args_ft)\n",
+    "print(\"\\n✅ Fine-tuning finished.\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# --- Step 4.3: Evaluate the Fine-Tuned Model ---\n",
+    "print(\"\\n\" + \"=\"*50)\n",
+    "print(\"\\n🚀 Starting inference on the BEST model from fine-tuning...\")\n",
+    "\n",
+    "best_model_path = f'./output_dir/{args_ft.task}/checkpoint-best.pth'\n",
+    "\n",
+    "if not os.path.exists(best_model_path):\n",
+    "    print(f\"❌ ERROR: The expected best model was not found at {best_model_path}\")\n",
+    "else:\n",
+    "    print(f\"Found best model at: {best_model_path}\")\n",
+    "    eval_parser = get_args_parser()\n",
+    "    args_eval_ft = eval_parser.parse_args([\n",
+    "        '--model', 'RETFound_mae',\n",
+    "        '--eval',\n",
+    "        '--data_path', '../IDRiD_data',\n",
+    "        '--resume', best_model_path\n",
+    "    ])\n",
+    "    finetune_main(args_eval_ft)\n",
+    "\n",
+    "print(\"\\n✅ Fine-tuned model evaluation finished.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# ==============================================================================\n",
+    "# ## 5. Workflow B: Evaluate a Pre-Trained Classifier\n",
+    "#\n",
+    "# This workflow allows you to skip fine-tuning and directly evaluate a\n",
+    "# model that has already been trained for this specific task.\n",
+    "# =============================================================================="
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# --- Step 5.1: Download Pre-Trained Classifier ---\n",
+    "checkpoint_path = './RETFound_IDRiD_Classifier.pth'\n",
+    "if not os.path.isfile(checkpoint_path):\n",
+    "    print(\"⏳ Model checkpoint not found. Downloading from Google Drive...\")\n",
+    "    # Note: Using the gdown ID for the classifier model\n",
+    "    !gdown --id 1b0grTwARX1cXnYnMB3ZJZES26aMkgkvZ -O {checkpoint_path}\n",
+    "    print(\"✅ Model download complete.\")\n",
+    "else:\n",
+    "    print(f\"✅ Pre-trained classifier already exists at '{checkpoint_path}'.\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# --- Step 5.2: Run Inference ---\n",
+    "print(\"\\n\" + \"=\"*50)\n",
+    "print(\"\\n🚀 Starting inference on the PRE-TRAINED classifier...\")\n",
+    "eval_parser_pretrained = get_args_parser()\n",
+    "args_eval_pretrained = eval_parser_pretrained.parse_args([\n",
+    "    '--model', 'RETFound_mae',\n",
+    "    '--eval',\n",
+    "    '--data_path', '../IDRiD_data',\n",
+    "    '--resume', checkpoint_path\n",
+    "])\n",
+    "finetune_main(args_eval_pretrained)\n",
+    "print(\"\\n✅ Pre-trained classifier evaluation finished.\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "colab": {
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}
diff --git a/engine_finetune.py b/engine_finetune.py
index fe60d44..0d7db6f 100644
--- a/engine_finetune.py
+++ b/engine_finetune.py
@@ -1,148 +1,60 @@
-import os
-import csv
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import numpy as np
-import matplotlib.pyplot as plt
-from typing import Iterable, Optional
-from timm.data import Mixup
-from timm.utils import accuracy
-from sklearn.metrics import (
-    accuracy_score, roc_auc_score, f1_score, average_precision_score,
-    hamming_loss, jaccard_score, recall_score, precision_score, cohen_kappa_score
-)
-from pycm import ConfusionMatrix
-import util.misc as misc
-import util.lr_sched as lr_sched
-
-def train_one_epoch(
-    model: torch.nn.Module,
-    criterion: torch.nn.Module,
-    data_loader: Iterable,
-    optimizer: torch.optim.Optimizer,
-    device: torch.device,
-    epoch: int,
-    loss_scaler,
-    max_norm: float = 0,
-    mixup_fn: Optional[Mixup] = None,
-    log_writer=None,
-    args=None
-):
-    """Train the model for one epoch."""
-    model.train(True)
-    metric_logger = misc.MetricLogger(delimiter="  ")
-    metric_logger.add_meter('lr', misc.SmoothedValue(window_size=1, fmt='{value:.6f}'))
-    print_freq, accum_iter = 20, args.accum_iter
-    optimizer.zero_grad()
-    
-    if log_writer:
-        print(f'log_dir: {log_writer.log_dir}')
-    
-    for data_iter_step, (samples, targets) in enumerate(metric_logger.log_every(data_loader, print_freq, f'Epoch: [{epoch}]')):
-        if data_iter_step % accum_iter == 0:
-            lr_sched.adjust_learning_rate(optimizer, data_iter_step / len(data_loader) + epoch, args)
-        
-        samples, targets = samples.to(device, non_blocking=True), targets.to(device, non_blocking=True)
-        if mixup_fn:
-            samples, targets = mixup_fn(samples, targets)
-        
-        with torch.cuda.amp.autocast():
-            outputs = model(samples)
-            loss = criterion(outputs, targets)
-        loss_value = loss.item()
-        loss /= accum_iter
-        
-        loss_scaler(loss, optimizer, clip_grad=max_norm, parameters=model.parameters(), create_graph=False,
-                    update_grad=(data_iter_step + 1) % accum_iter == 0)
-        if (data_iter_step + 1) % accum_iter == 0:
-            optimizer.zero_grad()
-        
-        torch.cuda.synchronize()
-        metric_logger.update(loss=loss_value)
-        min_lr = 10.
-        max_lr = 0.
-        for group in optimizer.param_groups:
-            min_lr = min(min_lr, group["lr"])
-            max_lr = max(max_lr, group["lr"])
-
-        metric_logger.update(lr=max_lr)
-
-        loss_value_reduce = misc.all_reduce_mean(loss_value)
-        if log_writer is not None and (data_iter_step + 1) % accum_iter == 0:
-            """ We use epoch_1000x as the x-axis in tensorboard.
-            This calibrates different curves when batch size changes.
-            """
-            epoch_1000x = int((data_iter_step / len(data_loader) + epoch) * 1000)
-            log_writer.add_scalar('loss/train', loss_value_reduce, epoch_1000x)
-            log_writer.add_scalar('lr', max_lr, epoch_1000x)
-    
-    metric_logger.synchronize_between_processes()
-    print("Averaged stats:", metric_logger)
-    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
-
-@torch.no_grad()
-def evaluate(data_loader, model, device, args, epoch, mode, num_class, log_writer):
-    """Evaluate the model."""
-    criterion = nn.CrossEntropyLoss()
-    metric_logger = misc.MetricLogger(delimiter="  ")
-    os.makedirs(os.path.join(args.output_dir, args.task), exist_ok=True)
-    
-    model.eval()
-    true_onehot, pred_onehot, true_labels, pred_labels, pred_softmax = [], [], [], [], []
-    
-    for batch in metric_logger.log_every(data_loader, 10, f'{mode}:'):
-        images, target = batch[0].to(device, non_blocking=True), batch[-1].to(device, non_blocking=True)
-        target_onehot = F.one_hot(target.to(torch.int64), num_classes=num_class)
-        
-        with torch.cuda.amp.autocast():
-            output = model(images)
-            loss = criterion(output, target)
-        output_ = nn.Softmax(dim=1)(output)
-        output_label = output_.argmax(dim=1)
-        output_onehot = F.one_hot(output_label.to(torch.int64), num_classes=num_class)
-        
-        metric_logger.update(loss=loss.item())
-        true_onehot.extend(target_onehot.cpu().numpy())
-        pred_onehot.extend(output_onehot.detach().cpu().numpy())
-        true_labels.extend(target.cpu().numpy())
-        pred_labels.extend(output_label.detach().cpu().numpy())
-        pred_softmax.extend(output_.detach().cpu().numpy())
-    
-    accuracy = accuracy_score(true_labels, pred_labels)
-    hamming = hamming_loss(true_onehot, pred_onehot)
-    jaccard = jaccard_score(true_onehot, pred_onehot, average='macro')
-    average_precision = average_precision_score(true_onehot, pred_softmax, average='macro')
-    kappa = cohen_kappa_score(true_labels, pred_labels)
-    f1 = f1_score(true_onehot, pred_onehot, zero_division=0, average='macro')
-    roc_auc = roc_auc_score(true_onehot, pred_softmax, multi_class='ovr', average='macro')
-    precision = precision_score(true_onehot, pred_onehot, zero_division=0, average='macro')
-    recall = recall_score(true_onehot, pred_onehot, zero_division=0, average='macro')
-    
-    score = (f1 + roc_auc + kappa) / 3
-    if log_writer:
-        for metric_name, value in zip(['accuracy', 'f1', 'roc_auc', 'hamming', 'jaccard', 'precision', 'recall', 'average_precision', 'kappa', 'score'],
-                                       [accuracy, f1, roc_auc, hamming, jaccard, precision, recall, average_precision, kappa, score]):
-            log_writer.add_scalar(f'perf/{metric_name}', value, epoch)
-    
-    print(f'val loss: {metric_logger.meters["loss"].global_avg}')
-    print(f'Accuracy: {accuracy:.4f}, F1 Score: {f1:.4f}, ROC AUC: {roc_auc:.4f}, Hamming Loss: {hamming:.4f},\n'
-          f' Jaccard Score: {jaccard:.4f}, Precision: {precision:.4f}, Recall: {recall:.4f},\n'
-          f' Average Precision: {average_precision:.4f}, Kappa: {kappa:.4f}, Score: {score:.4f}')
-    
-    metric_logger.synchronize_between_processes()
-    
-    results_path = os.path.join(args.output_dir, args.task, f'metrics_{mode}.csv')
-    file_exists = os.path.isfile(results_path)
-    with open(results_path, 'a', newline='', encoding='utf8') as cfa:
-        wf = csv.writer(cfa)
-        if not file_exists:
-            wf.writerow(['val_loss', 'accuracy', 'f1', 'roc_auc', 'hamming', 'jaccard', 'precision', 'recall', 'average_precision', 'kappa'])
-        wf.writerow([metric_logger.meters["loss"].global_avg, accuracy, f1, roc_auc, hamming, jaccard, precision, recall, average_precision, kappa])
-    
-    if mode == 'test':
-        cm = ConfusionMatrix(actual_vector=true_labels, predict_vector=pred_labels)
-        cm.plot(cmap=plt.cm.Blues, number_label=True, normalized=True, plot_lib="matplotlib")
-        plt.savefig(os.path.join(args.output_dir, args.task, 'confusion_matrix_test.jpg'), dpi=600, bbox_inches='tight')
-    
-    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}, score
+import math, sys
+from typing import Iterable, Optional
+import torch, torch.nn.functional as F
+import numpy as np
+from sklearn.metrics import roc_auc_score, classification_report, confusion_matrix
+from timm.utils import accuracy
+import util.misc as misc, util.lr_sched as lr_sched
+
+def train_one_epoch(model: torch.nn.Module, criterion: torch.nn.Module,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, loss_scaler, args=None):
+    model.train(True)
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', misc.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    header = f'Epoch: [{epoch}]'
+    for data_iter_step, (samples, targets) in enumerate(metric_logger.log_every(data_loader, 20, header)):
+        lr_sched.adjust_learning_rate(optimizer, data_iter_step / len(data_loader) + epoch, args)
+        samples, targets = samples.to(device, non_blocking=True), targets.to(device, non_blocking=True)
+        with torch.cuda.amp.autocast():
+            loss = criterion(model(samples), targets)
+        loss_value = loss.item()
+        if not math.isfinite(loss_value): sys.exit(f"Loss is {loss_value}, stopping training")
+        loss_scaler(loss, optimizer, parameters=model.parameters())
+        optimizer.zero_grad()
+        torch.cuda.synchronize()
+        metric_logger.update(loss=loss_value)
+        metric_logger.update(lr=optimizer.param_groups[0]["lr"])
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+@torch.no_grad()
+def evaluate(data_loader, model, device, args, epoch, mode, num_class, log_writer=None):
+    criterion = torch.nn.CrossEntropyLoss()
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    header, model.eval(), all_preds, all_labels, all_probs = 'Test:', [], [], []
+    for batch in metric_logger.log_every(data_loader, 10, header):
+        images, target = batch[0].to(device, non_blocking=True), batch[-1].to(device, non_blocking=True)
+        with torch.cuda.amp.autocast():
+            output = model(images)
+            loss = criterion(output, target)
+        preds, probs = torch.argmax(output, dim=1), F.softmax(output, dim=1)
+        all_preds.extend(preds.cpu().numpy())
+        all_labels.extend(target.cpu().numpy())
+        all_probs.extend(probs.cpu().detach().numpy())
+        metric_logger.update(loss=loss.item())
+        metric_logger.meters['acc1'].update(accuracy(output, target, topk=(1,))[0].item(), n=images.shape[0])
+    print(f'* Acc@1 {metric_logger.acc1.global_avg:.3f} loss {metric_logger.loss.global_avg:.3f}')
+    all_labels, all_preds, all_probs = np.array(all_labels), np.array(all_preds), np.array(all_probs)
+    print("\n--- Performance Metrics ---")
+    if len(np.unique(all_labels)) > 1 and len(np.unique(all_preds)) > 1:
+        try:
+            print(f"AUROC (Macro): {roc_auc_score(all_labels, all_probs, multi_class='ovr', average='macro'):.4f}")
+        except Exception as e: print(f"Could not calculate AUROC: {e}")
+    else: print("Skipping AUROC: not enough classes in labels or predictions.")
+    cm, fp = confusion_matrix(all_labels, all_preds), cm.sum(axis=0) - np.diag(cm)
+    tn = cm.sum() - (fp + (cm.sum(axis=1) - np.diag(cm)) + np.diag(cm))
+    print(f"Specificity (Macro): {np.mean(tn / (tn + fp)):.4f}")
+    print("\n--- Classification Report ---")
+    print(classification_report(all_labels, all_preds, target_names=[f'Class {i}' for i in range(num_class)], digits=4, zero_division=0))
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}, 'results'
diff --git a/main_finetune.py b/main_finetune.py
index 75c822d..49125f4 100644
--- a/main_finetune.py
+++ b/main_finetune.py
@@ -1,407 +1,123 @@
-import argparse
-import datetime
-import json
-
-import numpy as np
-import os
-import time
-from pathlib import Path
-
-import torch
-import torch.backends.cudnn as cudnn
-from torch.utils.tensorboard import SummaryWriter
-from timm.models.layers import trunc_normal_
-from timm.data.mixup import Mixup
-
-import models_vit as models
-import util.lr_decay as lrd
-import util.misc as misc
-from util.datasets import build_dataset
-from util.pos_embed import interpolate_pos_embed
-from util.misc import NativeScalerWithGradNormCount as NativeScaler
-from huggingface_hub import hf_hub_download, login
-from engine_finetune import train_one_epoch, evaluate
-
-import warnings
-import faulthandler
-
-faulthandler.enable()
-warnings.simplefilter(action='ignore', category=FutureWarning)
-
-
-def get_args_parser():
-    parser = argparse.ArgumentParser('MAE fine-tuning for image classification', add_help=False)
-    parser.add_argument('--batch_size', default=128, type=int,
-                        help='Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus')
-    parser.add_argument('--epochs', default=50, type=int)
-    parser.add_argument('--accum_iter', default=1, type=int,
-                        help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)')
-
-    # Model parameters
-    parser.add_argument('--model', default='vit_large_patch16', type=str, metavar='MODEL',
-                        help='Name of model to train')
-    parser.add_argument('--input_size', default=256, type=int,
-                        help='images input size')
-    parser.add_argument('--drop_path', type=float, default=0.2, metavar='PCT',
-                        help='Drop path rate (default: 0.1)')
-
-    # Optimizer parameters
-    parser.add_argument('--clip_grad', type=float, default=None, metavar='NORM',
-                        help='Clip gradient norm (default: None, no clipping)')
-    parser.add_argument('--weight_decay', type=float, default=0.05,
-                        help='weight decay (default: 0.05)')
-    parser.add_argument('--lr', type=float, default=None, metavar='LR',
-                        help='learning rate (absolute lr)')
-    parser.add_argument('--blr', type=float, default=5e-3, metavar='LR',
-                        help='base learning rate: absolute_lr = base_lr * total_batch_size / 256')
-    parser.add_argument('--layer_decay', type=float, default=0.65,
-                        help='layer-wise lr decay from ELECTRA/BEiT')
-    parser.add_argument('--min_lr', type=float, default=1e-6, metavar='LR',
-                        help='lower lr bound for cyclic schedulers that hit 0')
-    parser.add_argument('--warmup_epochs', type=int, default=10, metavar='N',
-                        help='epochs to warmup LR')
-
-    # Augmentation parameters
-    parser.add_argument('--color_jitter', type=float, default=None, metavar='PCT',
-                        help='Color jitter factor (enabled only when not using Auto/RandAug)')
-    parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1', metavar='NAME',
-                        help='Use AutoAugment policy. "v0" or "original". " + "(default: rand-m9-mstd0.5-inc1)'),
-    parser.add_argument('--smoothing', type=float, default=0.1,
-                        help='Label smoothing (default: 0.1)')
-
-    # * Random Erase params
-    parser.add_argument('--reprob', type=float, default=0.25, metavar='PCT',
-                        help='Random erase prob (default: 0.25)')
-    parser.add_argument('--remode', type=str, default='pixel',
-                        help='Random erase mode (default: "pixel")')
-    parser.add_argument('--recount', type=int, default=1,
-                        help='Random erase count (default: 1)')
-    parser.add_argument('--resplit', action='store_true', default=False,
-                        help='Do not random erase first (clean) augmentation split')
-
-    # * Mixup params
-    parser.add_argument('--mixup', type=float, default=0,
-                        help='mixup alpha, mixup enabled if > 0.')
-    parser.add_argument('--cutmix', type=float, default=0,
-                        help='cutmix alpha, cutmix enabled if > 0.')
-    parser.add_argument('--cutmix_minmax', type=float, nargs='+', default=None,
-                        help='cutmix min/max ratio, overrides alpha and enables cutmix if set (default: None)')
-    parser.add_argument('--mixup_prob', type=float, default=1.0,
-                        help='Probability of performing mixup or cutmix when either/both is enabled')
-    parser.add_argument('--mixup_switch_prob', type=float, default=0.5,
-                        help='Probability of switching to cutmix when both mixup and cutmix enabled')
-    parser.add_argument('--mixup_mode', type=str, default='batch',
-                        help='How to apply mixup/cutmix params. Per "batch", "pair", or "elem"')
-
-    # * Finetuning params
-    parser.add_argument('--finetune', default='', type=str,
-                        help='finetune from checkpoint')
-    parser.add_argument('--task', default='', type=str,
-                        help='finetune from checkpoint')
-    parser.add_argument('--global_pool', action='store_true')
-    parser.set_defaults(global_pool=True)
-    parser.add_argument('--cls_token', action='store_false', dest='global_pool',
-                        help='Use class token instead of global pool for classification')
-
-    # Dataset parameters
-    parser.add_argument('--data_path', default='./data/', type=str,
-                        help='dataset path')
-    parser.add_argument('--nb_classes', default=8, type=int,
-                        help='number of the classification types')
-    parser.add_argument('--output_dir', default='./output_dir',
-                        help='path where to save, empty for no saving')
-    parser.add_argument('--log_dir', default='./output_logs',
-                        help='path where to tensorboard log')
-    parser.add_argument('--device', default='cuda',
-                        help='device to use for training / testing')
-    parser.add_argument('--seed', default=0, type=int)
-    parser.add_argument('--resume', default='',
-                        help='resume from checkpoint')
-    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
-                        help='start epoch')
-    parser.add_argument('--eval', action='store_true',
-                        help='Perform evaluation only')
-    parser.add_argument('--dist_eval', action='store_true', default=False,
-                        help='Enabling distributed evaluation (recommended during training for faster monitor')
-    parser.add_argument('--num_workers', default=10, type=int)
-    parser.add_argument('--pin_mem', action='store_true',
-                        help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
-    parser.set_defaults(pin_mem=True)
-
-    # distributed training parameters
-    parser.add_argument('--world_size', default=1, type=int,
-                        help='number of distributed processes')
-    parser.add_argument('--local_rank', default=-1, type=int)
-    parser.add_argument('--dist_on_itp', action='store_true')
-    parser.add_argument('--dist_url', default='env://',
-                        help='url used to set up distributed training')
-
-    # fine-tuning parameters
-    parser.add_argument('--savemodel', action='store_true', default=True,
-                        help='Save model')
-    parser.add_argument('--norm', default='IMAGENET', type=str, help='Normalization method')
-    parser.add_argument('--enhance', action='store_true', default=False, help='Use enhanced data')
-    parser.add_argument('--datasets_seed', default=2026, type=int)
-
-    return parser
-
-
-def main(args, criterion):
-    if args.resume and not args.eval:
-        resume = args.resume
-        checkpoint = torch.load(args.resume, map_location='cpu')
-        print("Load checkpoint from: %s" % args.resume)
-        args = checkpoint['args']
-        args.resume = resume
-
-    misc.init_distributed_mode(args)
-
-    print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
-    print("{}".format(args).replace(', ', ',\n'))
-
-    device = torch.device(args.device)
-
-    # fix the seed for reproducibility
-    seed = args.seed + misc.get_rank()
-    torch.manual_seed(seed)
-    np.random.seed(seed)
-
-    cudnn.benchmark = True
-
-    if args.model=='RETFound_mae':
-        model = models.__dict__[args.model](
-        img_size=args.input_size,
-        num_classes=args.nb_classes,
-        drop_path_rate=args.drop_path,
-        global_pool=args.global_pool,
-    )
-    else:
-        model = models.__dict__[args.model](
-            num_classes=args.nb_classes,
-            drop_path_rate=args.drop_path,
-            args=args,
-        )
-    
-    if args.finetune and not args.eval:
-        
-        print(f"Downloading pre-trained weights from: {args.finetune}")
-        
-        checkpoint_path = hf_hub_download(
-            repo_id=f'YukunZhou/{args.finetune}',
-            filename=f'{args.finetune}.pth',
-        )
-        
-        checkpoint = torch.load(checkpoint_path, map_location='cpu')
-        print("Load pre-trained checkpoint from: %s" % args.finetune)
-        
-        if args.model!='RETFound_mae':
-            checkpoint_model = checkpoint['teacher']
-        else:
-            checkpoint_model = checkpoint['model']
-
-        checkpoint_model = {k.replace("backbone.", ""): v for k, v in checkpoint_model.items()}
-        checkpoint_model = {k.replace("mlp.w12.", "mlp.fc1."): v for k, v in checkpoint_model.items()}
-        checkpoint_model = {k.replace("mlp.w3.", "mlp.fc2."): v for k, v in checkpoint_model.items()}
-        
-        state_dict = model.state_dict()
-        for k in ['head.weight', 'head.bias']:
-            if k in checkpoint_model and checkpoint_model[k].shape != state_dict[k].shape:
-                print(f"Removing key {k} from pretrained checkpoint")
-                del checkpoint_model[k]
-
-        # interpolate position embedding
-        interpolate_pos_embed(model, checkpoint_model)
-
-        # load pre-trained model
-        msg = model.load_state_dict(checkpoint_model, strict=False)
-
-        trunc_normal_(model.head.weight, std=2e-5)
-
-    dataset_train = build_dataset(is_train='train', args=args)
-    dataset_val = build_dataset(is_train='val', args=args)
-    dataset_test = build_dataset(is_train='test', args=args)
-
-
-    if True:  # args.distributed:
-        num_tasks = misc.get_world_size()
-        global_rank = misc.get_rank()
-        if not args.eval:
-            sampler_train = torch.utils.data.DistributedSampler(
-                dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
-            )
-            print("Sampler_train = %s" % str(sampler_train))
-            if args.dist_eval:
-                if len(dataset_val) % num_tasks != 0:
-                    print(
-                        'Warning: Enabling distributed evaluation with an eval dataset not divisible by process number. '
-                        'This will slightly alter validation results as extra duplicate entries are added to achieve '
-                        'equal num of samples per-process.')
-                sampler_val = torch.utils.data.DistributedSampler(
-                    dataset_val, num_replicas=num_tasks, rank=global_rank,
-                    shuffle=True)  # shuffle=True to reduce monitor bias
-            else:
-                sampler_val = torch.utils.data.SequentialSampler(dataset_val)
-
-        if args.dist_eval:
-            if len(dataset_test) % num_tasks != 0:
-                print('Warning: Enabling distributed evaluation with an eval dataset not divisible by process number. '
-                      'This will slightly alter validation results as extra duplicate entries are added to achieve '
-                      'equal num of samples per-process.')
-            sampler_test = torch.utils.data.DistributedSampler(
-                dataset_test, num_replicas=num_tasks, rank=global_rank,
-                shuffle=True)  # shuffle=True to reduce monitor bias
-        else:
-            sampler_test = torch.utils.data.SequentialSampler(dataset_test)
-
-    if global_rank == 0 and args.log_dir is not None and not args.eval:
-        os.makedirs(args.log_dir, exist_ok=True)
-        log_writer = SummaryWriter(log_dir=os.path.join(args.log_dir, args.task))
-    else:
-        log_writer = None
-
-    if not args.eval:
-        data_loader_train = torch.utils.data.DataLoader(
-            dataset_train, sampler=sampler_train,
-            batch_size=args.batch_size,
-            num_workers=args.num_workers,
-            pin_memory=args.pin_mem,
-            drop_last=True,
-        )
-
-        print(f'len of train_set: {len(data_loader_train) * args.batch_size}')
-
-        data_loader_val = torch.utils.data.DataLoader(
-            dataset_val, sampler=sampler_val,
-            batch_size=args.batch_size,
-            num_workers=args.num_workers,
-            pin_memory=args.pin_mem,
-            drop_last=False
-        )
-
-    data_loader_test = torch.utils.data.DataLoader(
-        dataset_test, sampler=sampler_test,
-        batch_size=args.batch_size,
-        num_workers=args.num_workers,
-        pin_memory=args.pin_mem,
-        drop_last=False
-    )
-
-    mixup_fn = None
-    mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None
-    if mixup_active:
-        print("Mixup is activated!")
-        mixup_fn = Mixup(
-            mixup_alpha=args.mixup, cutmix_alpha=args.cutmix, cutmix_minmax=args.cutmix_minmax,
-            prob=args.mixup_prob, switch_prob=args.mixup_switch_prob, mode=args.mixup_mode,
-            label_smoothing=args.smoothing, num_classes=args.nb_classes)
-
-    if args.resume and args.eval:
-        checkpoint = torch.load(args.resume, map_location='cpu')
-        print("Load checkpoint from: %s" % args.resume)
-        model.load_state_dict(checkpoint['model'])
-
-    model.to(device)
-    model_without_ddp = model
-
-    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
-    print('number of model params (M): %.2f' % (n_parameters / 1.e6))
-
-    eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size()
-
-    if args.lr is None:  # only base_lr is specified
-        args.lr = args.blr * eff_batch_size / 256
-
-    print("base lr: %.2e" % (args.lr * 256 / eff_batch_size))
-    print("actual lr: %.2e" % args.lr)
-
-    print("accumulate grad iterations: %d" % args.accum_iter)
-    print("effective batch size: %d" % eff_batch_size)
-
-    if args.distributed:
-        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
-        model_without_ddp = model.module
-
-    no_weight_decay = model_without_ddp.no_weight_decay() if hasattr(model_without_ddp, 'no_weight_decay') else []
-    param_groups = lrd.param_groups_lrd(model_without_ddp, args.weight_decay,
-                                        no_weight_decay_list=no_weight_decay,
-                                        layer_decay=args.layer_decay
-                                        )
-    optimizer = torch.optim.AdamW(param_groups, lr=args.lr)
-    loss_scaler = NativeScaler()
-
-    print("criterion = %s" % str(criterion))
-
-    misc.load_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler)
-
-    if args.eval:
-        if 'epoch' in checkpoint:
-            print("Test with the best model at epoch = %d" % checkpoint['epoch'])
-        test_stats, auc_roc = evaluate(data_loader_test, model, device, args, epoch=0, mode='test',
-                                       num_class=args.nb_classes, log_writer=log_writer)
-        exit(0)
-
-    print(f"Start training for {args.epochs} epochs")
-    start_time = time.time()
-    max_score = 0.0
-    best_epoch = 0
-    for epoch in range(args.start_epoch, args.epochs):
-        if args.distributed:
-            data_loader_train.sampler.set_epoch(epoch)
-
-        train_stats = train_one_epoch(
-            model, criterion, data_loader_train,
-            optimizer, device, epoch, loss_scaler,
-            args.clip_grad, mixup_fn,
-            log_writer=log_writer,
-            args=args
-        )
-
-        val_stats, val_score = evaluate(data_loader_val, model, device, args, epoch, mode='val',
-                                        num_class=args.nb_classes, log_writer=log_writer)
-        if max_score < val_score:
-            max_score = val_score
-            best_epoch = epoch
-            if args.output_dir and args.savemodel:
-                misc.save_model(
-                    args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
-                    loss_scaler=loss_scaler, epoch=epoch, mode='best')
-        print("Best epoch = %d, Best score = %.4f" % (best_epoch, max_score))
-
-
-        if epoch == (args.epochs - 1):
-            checkpoint = torch.load(os.path.join(args.output_dir, args.task, 'checkpoint-best.pth'), map_location='cpu')
-            model_without_ddp.load_state_dict(checkpoint['model'], strict=False)
-            model.to(device)
-            print("Test with the best model, epoch = %d:" % checkpoint['epoch'])
-            test_stats, auc_roc = evaluate(data_loader_test, model, device, args, -1, mode='test',
-                                           num_class=args.nb_classes, log_writer=None)
-
-        if log_writer is not None:
-            log_writer.add_scalar('loss/val', val_stats['loss'], epoch)
-
-        log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
-                     'epoch': epoch,
-                     'n_parameters': n_parameters}
-
-        if args.output_dir and misc.is_main_process():
-            if log_writer is not None:
-                log_writer.flush()
-            with open(os.path.join(args.output_dir, args.task, "log.txt"), mode="a", encoding="utf-8") as f:
-                f.write(json.dumps(log_stats) + "\n")
-
-    total_time = time.time() - start_time
-    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
-    print('Training time {}'.format(total_time_str))
-
-
-if __name__ == '__main__':
-    args = get_args_parser()
-    args = args.parse_args()
-
-    criterion = torch.nn.CrossEntropyLoss()
-
-    if args.output_dir:
-        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
-    main(args, criterion)
-
-
+import argparse, datetime, json, numpy as np, os, time
+from pathlib import Path
+import torch, torch.backends.cudnn as cudnn
+from timm.data.mixup import Mixup
+import models_vit as models, util.lr_decay as lrd, util.misc as misc
+from util.datasets import build_dataset
+from util.misc import NativeScalerWithGradNormCount as NativeScaler
+from huggingface_hub import hf_hub_download
+from engine_finetune import train_one_epoch, evaluate
+import warnings
+
+warnings.simplefilter(action='ignore', category=FutureWarning)
+
+def get_args_parser():
+    parser = argparse.ArgumentParser('MAE fine-tuning for image classification', add_help=False)
+    parser.add_argument('--batch_size', default=128, type=int)
+    parser.add_argument('--epochs', default=50, type=int)
+    parser.add_argument('--accum_iter', default=1, type=int)
+    parser.add_argument('--model', default='vit_large_patch16', type=str)
+    parser.add_argument('--input_size', default=256, type=int)
+    parser.add_argument('--drop_path', type=float, default=0.2)
+    parser.add_argument('--clip_grad', type=float, default=None)
+    parser.add_argument('--weight_decay', type=float, default=0.05)
+    parser.add_argument('--lr', type=float, default=None)
+    parser.add_argument('--blr', type=float, default=5e-3)
+    parser.add_argument('--layer_decay', type=float, default=0.65)
+    parser.add_argument('--min_lr', type=float, default=1e-6)
+    parser.add_argument('--warmup_epochs', type=int, default=10)
+    parser.add_argument('--finetune', default='', type=str)
+    parser.add_argument('--task', default='', type=str)
+    parser.add_argument('--global_pool', action='store_true', default=True)
+    parser.add_argument('--data_path', default='./data/', type=str)
+    parser.add_argument('--nb_classes', default=5, type=int)
+    parser.add_argument('--output_dir', default='./output_dir')
+    parser.add_argument('--log_dir', default='./output_logs')
+    parser.add_argument('--device', default='cuda')
+    parser.add_argument('--seed', default=0, type=int)
+    parser.add_argument('--resume', default='')
+    parser.add_argument('--start_epoch', default=0, type=int)
+    parser.add_argument('--eval', action='store_true')
+    parser.add_argument('--num_workers', default=2, type=int)
+    parser.add_argument('--pin_mem', action='store_true', default=True)
+    parser.add_argument('--world_size', default=1, type=int)
+    parser.add_argument('--local_rank', default=-1, type=int)
+    parser.add_argument('--dist_on_itp', action='store_true')
+    parser.add_argument('--dist_url', default='env://')
+    return parser
+
+def main(args):
+    misc.init_distributed_mode(args)
+    device = torch.device(args.device)
+    seed = args.seed + misc.get_rank()
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    cudnn.benchmark = True
+
+    dataset_test = build_dataset(is_train='test', args=args)
+    if not args.eval:
+        dataset_train = build_dataset(is_train='train', args=args)
+        dataset_val = build_dataset(is_train='val', args=args)
+    else:
+        dataset_train, dataset_val = None, None
+
+    if args.distributed:
+        # ... (distributed setup omitted for Colab clarity)
+        pass
+    else:
+        sampler_test = torch.utils.data.SequentialSampler(dataset_test)
+        if not args.eval:
+            sampler_train = torch.utils.data.RandomSampler(dataset_train)
+            sampler_val = torch.utils.data.SequentialSampler(dataset_val)
+
+    log_writer = None # Disabled for Colab
+
+    data_loader_test = torch.utils.data.DataLoader(dataset_test, sampler=sampler_test, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False)
+    if not args.eval:
+        data_loader_train = torch.utils.data.DataLoader(dataset_train, sampler=sampler_train, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True)
+        data_loader_val = torch.utils.data.DataLoader(dataset_val, sampler=sampler_val, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False)
+
+    model = models.__dict__[args.model](img_size=args.input_size, num_classes=args.nb_classes, drop_path_rate=args.drop_path, global_pool=args.global_pool)
+
+    if args.finetune and not args.eval:
+        print(f"Downloading pre-trained weights from Hugging Face: {args.finetune}")
+        checkpoint_path = hf_hub_download(repo_id=f'YukunZhou/{args.finetune}', filename=f'{args.finetune}.pth')
+        checkpoint = torch.load(checkpoint_path, map_location='cpu')['model']
+        msg = model.load_state_dict(checkpoint, strict=False)
+        print(f"Loaded pre-trained checkpoint from {args.finetune} with message: {msg}")
+
+    if args.resume:
+        print(f"Resuming from checkpoint: {args.resume}")
+        checkpoint = torch.load(args.resume, map_location='cpu')['model']
+        model.load_state_dict(checkpoint, strict=False)
+
+    model.to(device)
+    print(f'Number of model params (M): {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1.e6:.2f}')
+
+    if args.eval:
+        evaluate(data_loader_test, model, device, args, 0, 'test', args.nb_classes, log_writer)
+        return
+
+    eff_batch_size = args.batch_size * misc.get_world_size()
+    if args.lr is None: args.lr = args.blr * eff_batch_size / 256
+    print(f"Actual lr: {args.lr:.2e}")
+    param_groups = lrd.param_groups_lrd(model, args.weight_decay, no_weight_decay_list=model.no_weight_decay(), layer_decay=args.layer_decay)
+    optimizer = torch.optim.AdamW(param_groups, lr=args.lr)
+    loss_scaler = NativeScaler()
+    criterion = torch.nn.CrossEntropyLoss()
+
+    print(f"--- Starting Training for {args.epochs} epochs ---")
+    max_accuracy = 0.0
+    for epoch in range(args.start_epoch, args.epochs):
+        train_stats = train_one_epoch(model, criterion, data_loader_train, optimizer, device, epoch, loss_scaler, args=args)
+        val_stats, _ = evaluate(data_loader_val, model, device, args, epoch, 'val', args.nb_classes, log_writer=log_writer)
+        print(f"EPOCH:{epoch} | Val Acc: {val_stats['acc1']:.1f}%")
+        if max_accuracy < val_stats["acc1"]:
+            max_accuracy = val_stats["acc1"]
+            misc.save_model(args=args, model=model, model_without_ddp=model, optimizer=optimizer, loss_scaler=loss_scaler, epoch=epoch, mode='best')
+        print(f'Max accuracy: {max_accuracy:.2f}%')
+
+if __name__ == '__main__':
+    args = get_args_parser().parse_args()
+    if args.output_dir: Path(os.path.join(args.output_dir, args.task)).mkdir(parents=True, exist_ok=True)
+    main(args)