File size: 4,440 Bytes
53e66de | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | {
"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
}
|