Example_Energy_Based_OOD_Detection.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "authorship_tag": "ABX9TyNhEhbm7FS5Jz3vKzMHfGQU",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/OneFineStarstuff/OneFineStardust/blob/main/Example_Energy_Based_OOD_Detection.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "toXeUIBf_vN-"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from transformers import LongformerForSequenceClassification, LongformerTokenizer\n",
        "\n",
        "class EnergyBasedOOD:\n",
        "    def __init__(self, model):\n",
        "        self.model = model\n",
        "\n",
        "    def compute_energy(self, logits):\n",
        "        return -torch.logsumexp(logits, dim=1)\n",
        "\n",
        "    def detect_ood(self, input_ids, attention_mask, threshold):\n",
        "        # Ensure the model is in evaluation mode\n",
        "        self.model.eval()\n",
        "\n",
        "        # Disable gradient computation for inference\n",
        "        with torch.no_grad():\n",
        "            outputs = self.model(input_ids, attention_mask=attention_mask)\n",
        "            logits = outputs.logits\n",
        "\n",
        "        # Compute the energy scores\n",
        "        energy = self.compute_energy(logits)\n",
        "\n",
        "        # Return OOD detection results\n",
        "        return energy > threshold\n",
        "\n",
        "# Example usage\n",
        "# Initialize tokenizer\n",
        "tokenizer = LongformerTokenizer.from_pretrained(\"allenai/longformer-base-4096\")\n",
        "\n",
        "# Initialize the model for sequence classification\n",
        "model = LongformerForSequenceClassification.from_pretrained(\"allenai/longformer-base-4096\", num_labels=2)\n",
        "\n",
        "# Initialize the OOD detector\n",
        "ood_detector = EnergyBasedOOD(model)\n",
        "\n",
        "# Dummy input data (Ensure this is a tensor of the correct shape)\n",
        "text_data = [\"This is a sample text for model explanation.\"]\n",
        "inputs = tokenizer(text_data, return_tensors=\"pt\", truncation=True, padding=True, max_length=64)\n",
        "input_ids = inputs[\"input_ids\"]\n",
        "attention_mask = inputs[\"attention_mask\"]\n",
        "\n",
        "# Define a threshold for OOD detection\n",
        "threshold = 0.5\n",
        "\n",
        "# Detect OOD\n",
        "ood_results = ood_detector.detect_ood(input_ids, attention_mask, threshold)\n",
        "print(\"OOD Detection Results:\", ood_results)"
      ]
    }
  ]
}