data_protein_esm.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import os\n",
    "import pandas as pd\n",
    "import json\n",
    "dataset = 'davis'\n",
    "\n",
    "test_data = pd.read_csv('data/'+dataset+'/'+dataset+'_test.csv')\n",
    "train_data = pd.read_csv('data/'+dataset+'/'+dataset+'_train.csv')\n",
    "\n",
    "complete_data=pd.concat([train_data,test_data],axis=0)\n",
    "\n",
    "protein_set = list(set(complete_data['target_sequence']))\n",
    "drug_set=list(set(complete_data['compound_iso_smiles']))\n",
    "print(len(protein_set))\n",
    "print(len(drug_set))\n",
    "\n",
    "protein_dict= {}\n",
    "for i in range(len(protein_set)):\n",
    "    protein_dict[i] = protein_set[i]\n",
    "with open('data/'+dataset+'/protein_dict.txt','w') as file:\n",
    "    file.write(json.dumps(protein_dict))\n",
    "\n",
    "drug_dict= {}\n",
    "for i in range(len(drug_set)):\n",
    "    drug_dict[i] = drug_set[i]\n",
    "with open('data/'+dataset+'/drug_dict.txt','w') as file:\n",
    "    file.write(json.dumps(drug_dict))\n",
    "\n",
    "result=''\n",
    "for key, value in protein_dict.items():\n",
    "    result+='>UniRef50_'+str(key)+'\\n'+str(value)+'\\n'\n",
    "f = open('data/'+dataset+'/proteins.fasta', 'w')\n",
    "f.write(result)\n",
    "\n",
    "if not os.path.isdir('data/'+dataset+'/'+'/proteins_emb_esm2'):\n",
    "    os.mkdir('data/'+dataset+'/'+'/protein_emb_esm2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python scripts/extract.py esm2_t36_3B_UR50D data/davis/proteins.fasta data/davis/proteins_emb_esm2 --include mean per_tok"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}