{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":87793,"databundleVersionId":12024591,"sourceType":"competition"}],"dockerImageVersionId":30918,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Compute Embeddings using Helix model from HelicalAI\n\nThis notebook uses the Helix-mRNA model from HelicalAI (https://arxiv.org/abs/2502.13785) to compute embeddings from the training sequences. This model is trained on a large dataset of sequences, and it is able to capture features important for RNA structure. The embeddings can then be used for further training.\n\nWORKING VERSION\n\noriginal version: https://www.kaggle.com/code/dalloliogm/computing-embeddings-using-helix-mrna","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19"}},{"cell_type":"markdown","source":"## Install libraries\n\nInstalling libraries is complicated, because there are many dependencies, and some rely on older versions of pandas and other packages. ","metadata":{}},{"cell_type":"code","source":"!pip uninstall -y cupy -q\n!pip uninstall -y cupy-cuda12x -q\n!pip install cupy-cuda11x -q\n\n# https://claude.ai/chat/ca65c790-2155-4ce0-8f7b-d47741ec4f62\n\n# Make sure torch and torchvision versions match\n!pip install torch==2.0.1 torchvision==0.15.2 -q\n!pip install helical -q","metadata":{"trusted":true,"_kg_hide-output":true,"_kg_hide-input":true,"scrolled":true,"execution":{"iopub.status.busy":"2025-05-03T04:08:11.566053Z","iopub.execute_input":"2025-05-03T04:08:11.566277Z","iopub.status.idle":"2025-05-03T04:13:28.679430Z","shell.execute_reply.started":"2025-05-03T04:08:11.566257Z","shell.execute_reply":"2025-05-03T04:13:28.678340Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Try installing helical normally instead of to a custom location\n!pip install helical -q --upgrade\n#!pip install helical -q # --target=/kaggle/working/ --upgrade #--use-deprecated=legacy-resolver\nfrom helical.models.helix_mrna import HelixmRNAConfig, HelixmRNA, HelixmRNAFineTuningModel\n","metadata":{"trusted":true,"_kg_hide-input":false,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2025-05-03T04:13:28.680431Z","iopub.execute_input":"2025-05-03T04:13:28.680681Z","iopub.status.idle":"2025-05-03T04:13:52.897129Z","shell.execute_reply.started":"2025-05-03T04:13:28.680661Z","shell.execute_reply":"2025-05-03T04:13:52.896423Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import subprocess\nimport torch\nimport pandas as pd","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-03T04:13:52.898223Z","iopub.execute_input":"2025-05-03T04:13:52.898966Z","iopub.status.idle":"2025-05-03T04:13:52.902738Z","shell.execute_reply.started":"2025-05-03T04:13:52.898918Z","shell.execute_reply":"2025-05-03T04:13:52.901677Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Read Data","metadata":{}},{"cell_type":"code","source":"# Load data\n\ntrain_sequences=pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/train_sequences.csv\")\ntrain_labels=pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/train_labels.csv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-03T04:13:52.904624Z","iopub.execute_input":"2025-05-03T04:13:52.904816Z","iopub.status.idle":"2025-05-03T04:13:53.588006Z","shell.execute_reply.started":"2025-05-03T04:13:52.904799Z","shell.execute_reply":"2025-05-03T04:13:53.587102Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Import Helix and compute embeddings","metadata":{}},{"cell_type":"code","source":"from helical.models.helix_mrna import HelixmRNAConfig, HelixmRNA, HelixmRNAFineTuningModel\n\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-03T04:13:53.589067Z","iopub.execute_input":"2025-05-03T04:13:53.589316Z","iopub.status.idle":"2025-05-03T04:13:53.593077Z","shell.execute_reply.started":"2025-05-03T04:13:53.589297Z","shell.execute_reply":"2025-05-03T04:13:53.592167Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# We set the max length to the maximum length of the sequences in the training data + 10 to include space for special tokens\nhelix_mrna_config = HelixmRNAConfig(device=device, batch_size=1, max_length=max(len(s) for s in train_sequences[\"sequence\"])+10)\nhelix_mrna = HelixmRNA(helix_mrna_config)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-03T04:13:53.593880Z","iopub.execute_input":"2025-05-03T04:13:53.594143Z","iopub.status.idle":"2025-05-03T04:13:56.207174Z","shell.execute_reply.started":"2025-05-03T04:13:53.594113Z","shell.execute_reply":"2025-05-03T04:13:56.206483Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def replace_chars(text, char_list, replacement):\n    for char in char_list:\n        text = text.replace(char, replacement)\n    return text","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-03T04:13:56.207996Z","iopub.execute_input":"2025-05-03T04:13:56.208302Z","iopub.status.idle":"2025-05-03T04:13:56.212185Z","shell.execute_reply.started":"2025-05-03T04:13:56.208270Z","shell.execute_reply":"2025-05-03T04:13:56.211411Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"replacement = \"\"\nchars_to_replace = [\"X\", \"-\"]\n\n#processed_train_data = helix_mrna.process_data([s.replace(\"X\", \"\") for s in train_sequences['sequence']])\nprocessed_train_data = helix_mrna.process_data([replace_chars(s, chars_to_replace, replacement) for s in train_sequences['sequence']])\n\nembeddings = helix_mrna.get_embeddings(processed_train_data)\nembeddings = embeddings[:, -2, :]\nprint(embeddings.shape)\nprint(embeddings[:1][0:10])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-03T04:13:56.212845Z","iopub.execute_input":"2025-05-03T04:13:56.213105Z","iopub.status.idle":"2025-05-03T04:21:47.222601Z","shell.execute_reply.started":"2025-05-03T04:13:56.213085Z","shell.execute_reply":"2025-05-03T04:21:47.221511Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Save embeddings","metadata":{}},{"cell_type":"code","source":"df = pd.DataFrame(embeddings)\ndf.to_csv(\"/kaggle/working/train_helix_embeddings.csv\", index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-03T04:21:47.223597Z","iopub.execute_input":"2025-05-03T04:21:47.223841Z","iopub.status.idle":"2025-05-03T04:21:47.440151Z","shell.execute_reply.started":"2025-05-03T04:21:47.223819Z","shell.execute_reply":"2025-05-03T04:21:47.439272Z"}},"outputs":[],"execution_count":null}]}