{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "from tqdm import tqdm\n",
    "\n",
    "import pandas as pd\n",
    "import torch\n",
    "from transformers import AutoTokenizer, BigBirdForMaskedLM\n",
    "\n",
    "from CodonTransformer.CodonJupyter import (\n",
    "    UserContainer,\n",
    "    display_organism_dropdown,\n",
    "    display_protein_input,\n",
    "    format_model_output,\n",
    ")\n",
    "from CodonTransformer.CodonPrediction import predict_dna_sequence\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load model and tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"adibvafa/CodonTransformer\")\n",
    "model = BigBirdForMaskedLM.from_pretrained(\"adibvafa/CodonTransformer\").to(DEVICE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Optimizing a Single Sequence**\n",
    "-------------------------------------\n",
    "1. Run the next code cell and input only your protein sequence and organism\n",
    "\n",
    "2. Run the code cell after it to optimize the sequence and display it.\n",
    "\n",
    "Protein sequences should end with \"*\" or \"_\" or an amino acid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sample: MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKTRREAEDLQVGQVELGG, Homo sapiens\n",
    "user = UserContainer()\n",
    "display_protein_input(user)\n",
    "display_organism_dropdown(user)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = predict_dna_sequence(\n",
    "    protein=user.protein,\n",
    "    organism=user.organism,\n",
    "    device=DEVICE,\n",
    "    tokenizer=tokenizer,\n",
    "    model=model,\n",
    "    attention_type=\"original_full\",\n",
    "    deterministic=True,\n",
    "    # Can set temperature for non deterministic prediction\n",
    ")\n",
    "\n",
    "print(format_model_output(output))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Optimizing Multiple Sequences**\n",
    "-------------------------------------\n",
    "1. Create a CSV file that has columns 'protein_sequence' and 'organism'.\n",
    "   You can have other columns in any order.\n",
    "\n",
    "2. Replace the _dataset_path_ below with the actual path to your CSV file.\n",
    "\n",
    "3. Run the next code cells to optimize and save the predicted DNA sequences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Update with the actual path to your dataset\n",
    "dataset_path = \"scripts/demo/sample_dataset.csv\"\n",
    "output_path = \"scripts/demo/sample_predictions.csv\"\n",
    "\n",
    "dataset = pd.read_csv(dataset_path, index_col=0)\n",
    "dataset[\"predicted_dna\"] = None\n",
    "dataset.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for index, data in tqdm(\n",
    "    dataset.iterrows(),\n",
    "    desc=f\"CodonTransformer Predicting\",\n",
    "    unit=\" Sequences\",\n",
    "    total=dataset.shape[0],\n",
    "):\n",
    "\n",
    "    outputs = predict_dna_sequence(\n",
    "        protein=data[\"protein_sequence\"],\n",
    "        organism=data[\"organism\"],\n",
    "        device=DEVICE,\n",
    "        tokenizer_object=tokenizer,\n",
    "        model_object=model,\n",
    "    )\n",
    "    dataset.loc[index, \"predicted_dna\"] = outputs.predicted_dna\n",
    "\n",
    "dataset.to_csv(output_path)\n",
    "dataset.head()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "light",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}