notebooks/model_notebooks/AntiBERTy.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "40cc2457-1ccc-420f-bd98-ee6f196849fd",
   "metadata": {},
   "source": [
    "## AntiBERTy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb236a3e-35bc-40da-bad4-62173307320d",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72920c06-cfff-4263-b627-1a584122992f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "from prtm.models.antiberty import config\n",
    "from prtm.models.antiberty.modeling import (\n",
    "    AntiBERTyForSequenceEmbedding, \n",
    "    AntiBERTyForSequenceInpainting, \n",
    "    AntiBERTyForAntibodySequenceClassification,\n",
    ")\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dc82aa4-04e2-4175-a3ed-270f3d387298",
   "metadata": {},
   "outputs": [],
   "source": [
    "sequence = \"EVQLVQSGPEVKKPGTSVKVSCKASGFTFMSSAVQWVRQARGQRLEWIGWIVIGSGNTNYAQKFQERVTITRDMSTSTAYMELSSLRSEDTAVYYCAAPYCSSISCNDGFDIWGQGTMVTVS\"\n",
    "masked_sequence = \"\".join([c if random.random() > 0.2 else \"_\" for c in sequence])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f191f29-a2dd-47c2-a325-50cfef08778f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sequence_embedder = AntiBERTyForSequenceEmbedding()\n",
    "sequence_inpainter = AntiBERTyForSequenceInpainting()\n",
    "antibody_classifier = AntiBERTyForAntibodySequenceClassification()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d62e8fdf-adcf-40e7-a374-6d03717f9590",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding, aux_output = sequence_embedder(sequence)\n",
    "inpainted_sequence, _ = sequence_inpainter(masked_sequence)\n",
    "antibody_classes, _ = antibody_classifier(sequence)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7e0a1c7-c4e6-41b0-a365-be45c04cedde",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Masked sequence:    \", masked_sequence)\n",
    "print(\"Inpainted sequence: \", inpainted_sequence)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31eb8d5f-ff54-4cb2-931c-43143df9dd3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Antibody predicted classes:\", antibody_classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31b5bd36-9b2c-47f1-a3f8-4b9e37eeeab4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot one of the 64 residue attention maps\n",
    "plt.imshow(aux_output[\"attention\"][0].cpu().numpy(), cmap=\"jet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e993715-09a7-4bad-a358-3f552b9ac142",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}