{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":87793,"databundleVersionId":11553390,"sourceType":"competition"}],"dockerImageVersionId":30918,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Computing 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\nNote: I'm struggling with installing helical on kaggle notebooks. This version may not work.","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":"import subprocess\nimport torch\nimport pandas as pd","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-25T13:48:26.026516Z","iopub.execute_input":"2025-03-25T13:48:26.026859Z","iopub.status.idle":"2025-03-25T13:48:29.652408Z","shell.execute_reply.started":"2025-03-25T13:48:26.026823Z","shell.execute_reply":"2025-03-25T13:48:29.651557Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install helical -q --target=/kaggle/working/ --upgrade #--use-deprecated=legacy-resolver\nfrom helical.models.helix_mrna import HelixmRNAConfig, HelixmRNA, HelixmRNAFineTuningModel","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-25T13:48:29.653358Z","iopub.execute_input":"2025-03-25T13:48:29.653877Z","iopub.status.idle":"2025-03-25T13:52:24.687821Z","shell.execute_reply.started":"2025-03-25T13:48:29.653842Z","shell.execute_reply":"2025-03-25T13:52:24.685700Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip uninstall -y cupy -q\n!pip uninstall -y cupy-cuda12x -q\n!pip install cupy-cuda11x -q\n!pip uninstall -y torch torchvision torchaudio\n!pip install -q torch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1\n\n#pip install helical -q --target=/kaggle/working/ --upgrade\n!pip install -q --upgrade git+https://github.com/helicalAI/helical.git\nfrom helical.models.helix_mrna import HelixmRNAConfig, HelixmRNA, HelixmRNAFineTuningModel\n","metadata":{"trusted":true,"_kg_hide-output":true,"_kg_hide-input":true,"scrolled":true,"execution":{"iopub.status.busy":"2025-03-25T13:52:24.688290Z","iopub.status.idle":"2025-03-25T13:52:24.688562Z","shell.execute_reply":"2025-03-25T13:52:24.688453Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\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-03-25T13:52:24.689496Z","iopub.status.idle":"2025-03-25T13:52:24.689889Z","shell.execute_reply":"2025-03-25T13:52:24.689734Z"}},"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-03-25T13:52:24.690658Z","iopub.status.idle":"2025-03-25T13:52:24.691035Z","shell.execute_reply":"2025-03-25T13:52:24.690877Z"}},"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-03-25T13:52:24.691903Z","iopub.status.idle":"2025-03-25T13:52:24.692273Z","shell.execute_reply":"2025-03-25T13:52:24.692115Z"}},"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)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-25T13:52:24.693140Z","iopub.status.idle":"2025-03-25T13:52:24.693499Z","shell.execute_reply":"2025-03-25T13:52:24.693344Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"processed_train_data = helix_mrna.process_data([s.replace(\"X\", \"\") for s in train_sequences['sequence']])\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-03-25T13:52:24.694722Z","iopub.status.idle":"2025-03-25T13:52:24.695078Z","shell.execute_reply":"2025-03-25T13:52:24.694922Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"embeddings.to_csv(\"/kaggle/output/train_helix_embeddings.csv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-25T13:52:24.696123Z","iopub.status.idle":"2025-03-25T13:52:24.696506Z","shell.execute_reply":"2025-03-25T13:52:24.696336Z"}},"outputs":[],"execution_count":null}]}