Download data files and minor refactoring

.gitignore

@@ -1,2 +1,6 @@
 .DS_Store
 .ipynb_checkpoints
+venv/
+__pycache__
+data/
+.env

gene_embeddings_examples.ipynb

@@ -10,8 +10,10 @@
     "import json\n",
     "import numpy as np\n",
     "import openai\n",
-    "# remember to set your OpenAI api key\n",
-    "openai.api_key = ''"
+    "\n",
+    "# Add your OpenAI API key to the .env file as OPENAI_API_KEY=your_api_key\n",
+    "import dotenv\n",
+    "dotenv.load_dotenv()\n"
    ]
   },
   {

request_ncbi_text_for_genes.ipynb

@@ -2,7 +2,39 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
+   "id": "68ec84f8",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Collecting html2text\n",
+      "  Downloading html2text-2024.2.26.tar.gz (56 kB)\n",
+      "  Installing build dependencies ... \u001b[?25ldone\n",
+      "\u001b[?25h  Getting requirements to build wheel ... \u001b[?25ldone\n",
+      "\u001b[?25h  Preparing metadata (pyproject.toml) ... \u001b[?25ldone\n",
+      "\u001b[?25hBuilding wheels for collected packages: html2text\n",
+      "  Building wheel for html2text (pyproject.toml) ... \u001b[?25ldone\n",
+      "\u001b[?25h  Created wheel for html2text: filename=html2text-2024.2.26-py3-none-any.whl size=33111 sha256=af187fe8ca5aa9aba20968445ca5ffb234f15e7db400f944bfeb37a7fb5cbbaf\n",
+      "  Stored in directory: /Users/rj/Library/Caches/pip/wheels/2b/01/23/578505d65e2a97d78bf1fe3fc8256ecf37572dc1df598b0eaf\n",
+      "Successfully built html2text\n",
+      "Installing collected packages: html2text\n",
+      "Successfully installed html2text-2024.2.26\n",
+      "\n",
+      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.3.1\u001b[0m\n",
+      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "!pip install html2text"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
    "id": "5e33a685",
    "metadata": {},
    "outputs": [],
@@ -110,18 +142,6 @@
     "gene_name_to_tax_id"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "2b2a50fd",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "with open('/Users/yiquntchen/Downloads/human/vocab.json', 'rb') as handle:\n",
-    "    vocab_gene = json.load(handle)\n",
-    "vocab_gene_list = list(vocab_gene.keys())"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -168,18 +188,62 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d609d3ad",
+   "id": "b7d67f42",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# load genes used in GenePT\n",
-    "with open(f\"{data_dir}/vocab.json', 'rb') as handle:\n",
-    "    vocab_gene = json.load(handle)\n",
-    "vocab_gene_list = list(vocab_gene.keys())\n",
     "\n",
-    "# load genes used in Geneformer\n",
-    "with open(f\"{data_dir}/token_dictionary.pkl\", 'rb') as handle:\n",
-    "    token_dictionary = pickle.load(handle)\n",
+    "from src.utils import download_gdrive_file\n",
+    "\n",
+    "scgpt_vocab_url = \"https://drive.google.com/file/d/1H3E_MJ-Dl36AQV6jLbna2EdvgPaqvqcC/view?usp=drive_link\"\n",
+    "scgpt_vocab_path = \"data/vocab.json\"\n",
+    "download_gdrive_file(scgpt_vocab_url, scgpt_vocab_path)\n",
+    "vocab_gene = json.load(open(scgpt_vocab_path, 'rb'))\n",
+    "vocab_gene_list = list(vocab_gene.keys())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "c7f5f18a",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "File already exists at data/token_dictionary.pkl\n"
+     ]
+    }
+   ],
+   "source": [
+    "from src.utils import download_file\n",
+    "import pickle\n",
+    "\n",
+    "geneformer_token_dict_url = \"https://huggingface.co/ctheodoris/Geneformer/resolve/main/geneformer/token_dictionary_gc95M.pkl\"\n",
+    "geneformer_token_dict_path = \"data/token_dictionary.pkl\"\n",
+    "download_file(geneformer_token_dict_url, geneformer_token_dict_path)\n",
+    "token_dictionary = pickle.load(open(geneformer_token_dict_path, 'rb'))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "d609d3ad",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "6975 input query terms found dup hits:\t[('A2ML1-AS1', 2), ('A2ML1-AS2', 2), ('A2MP1', 2), ('AACSP1', 2), ('AADACL2-AS1', 3), ('AADACP1', 2)\n",
+      "21975 input query terms found no hit:\t['5S_rRNA_ENSG00000276861', '5S_rRNA_ENSG00000277411', '5S_rRNA_ENSG00000277488', '5S_rRNA_ENSG00000\n",
+      "7 input query terms found dup hits:\t[('ENSG00000267635', 2), ('ENSG00000268674', 3), ('ENSG00000275921', 2), ('ENSG00000278277', 3), ('E\n",
+      "11 input query terms found no hit:\t['<cls>', '<eos>', '<mask>', '<pad>', 'ENSG00000139656', 'ENSG00000168078', 'ENSG00000189144', 'ENSG\n"
+     ]
+    }
+   ],
+   "source": [
     "\n",
     "# example query to convert gene IDs into page ids for NCBI \n",
     "vocab_gene_list_results = mg.querymany(sorted(vocab_gene_list), scopes='symbol', species='human')\n",
@@ -203,7 +267,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.4"
+   "version": "3.12.7"
   }
  },
  "nbformat": 4,

src/utils.py

@@ -0,0 +1,154 @@
+from pathlib import Path
+import pickle
+import requests
+import zipfile
+from shutil import move
+
+def download_file(url: str, output_path: Path | str | None = None, chunk_size: int = 8192) -> Path:
+    """
+    Download a file in chunks and return the Path to the downloaded file.
+    If output_path is not provided, the filename will be extracted from the URL.
+
+    Args:
+        url: URL of the file to download
+        output_path: Path where the file should be saved (optional)
+        chunk_size: Size of chunks to download (default: 8192 bytes)
+
+    Returns:
+        Path: Path object pointing to the downloaded file
+
+    Raises:
+        requests.exceptions.RequestException: If download fails
+        ValueError: If URL doesn't contain a filename and output_path is not provided
+    """
+    # If no output path is provided, extract filename from URL
+    if output_path is None:
+        # Get the filename from the URL
+        filename = url.split('/')[-1].split('?')[0]  # Remove query parameters
+        if not filename:
+            raise ValueError("Could not determine filename from URL. Please provide output_path.")
+        output_path = Path(__file__).parent.parent / "data" / filename
+    else:
+        output_path = Path(output_path)
+
+    # Create parent directories if they don't exist
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Download the file if it doesn't exist
+    if not output_path.exists():
+        print(f"Downloading file to {output_path}...")
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
+
+        with open(output_path, "wb") as f:
+            for chunk in response.iter_content(chunk_size=chunk_size):
+                f.write(chunk)
+        print("Download complete!")
+    else:
+        print(f"File already exists at {output_path}")
+
+    return output_path
+
+
+
+def extract_zip(zip_path: Path | str, extract_dir: Path | str) -> None:
+    """
+    Extract a zip file, skipping files that already exist with the same size.
+
+    Args:
+        zip_path: Path to the zip file
+        extract_dir: Directory where files should be extracted
+
+    Returns:
+        None
+    """
+    # Convert to Path objects
+    zip_path = Path(zip_path)
+    extract_dir = Path(extract_dir)
+
+    # Create extraction directory if it doesn't exist
+    extract_dir.mkdir(parents=True, exist_ok=True)
+
+    print("Extracting files...")
+    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+        # Get list of files in archive
+        for file_info in zip_ref.infolist():
+            target_path = extract_dir / file_info.filename
+
+            # Check if file needs to be extracted
+            extract_file = True
+            if target_path.exists():
+                # Compare file sizes
+                if target_path.stat().st_size == file_info.file_size:
+                    extract_file = False
+                    print(f"Skipping {file_info.filename} - already exists with same size")
+
+            if extract_file:
+                print(f"Extracting {file_info.filename}")
+                zip_ref.extract(file_info, extract_dir)
+
+    print("Extraction complete!")
+
+def setup_data_dir():
+    """
+    Set up the data directory by downloading and extracting required files.
+
+    Downloads the GenePT embedding file from Zenodo and extracts it into
+    the 'data' directory. If files already exist, downloading and extraction
+    will be skipped.
+    """
+
+    # Create data directory if it doesn't exist
+    data_dir = Path("data")
+    data_dir.mkdir(exist_ok=True)
+
+    # Download URL
+    url = "https://zenodo.org/records/10833191/files/GenePT_emebdding_v2.zip?download=1"
+    zip_path = download_file(url)
+
+    extract_zip(zip_path, data_dir)
+
+    print("Setup finished!")
+
+def get_gene_embeddings(model_name) -> dict:
+    """
+    Get the GenePT embeddings for a given model.
+    """
+    model_path_map = {
+        "text-embedding-ada-002": "data/GenePT_emebdding_v2/GenePT_gene_embedding_ada_text.pickle",
+        "text-embedding-3-large": "data/GenePT_emebdding_v2/GenePT_gene_protein_embedding_model_3_text.pickle.",
+    }
+
+    model_path = Path(__file__).parent.parent / model_path_map[model_name]
+    with open(model_path, "rb") as f:
+        return pickle.load(f)
+
+def download_gdrive_file(url: str, output_path: Path | str ) -> Path:
+    """
+    Download a file from Google Drive using gdown and optionally rename/move it.
+
+    Args:
+        url: Google Drive sharing URL
+        output_path: Path where the file should be saved (optional)
+
+    Returns:
+        Path: Path object pointing to the downloaded file
+    """
+    import gdown
+
+    output_path = Path(output_path)
+    if not output_path.exists():
+
+        # Download to current directory first
+        downloaded_path = gdown.download(url, fuzzy=True)
+        if not downloaded_path:
+            raise RuntimeError("Download failed")
+        temp_path = Path(downloaded_path)    
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Move file to final location if it's different from where it was downloaded
+        if temp_path != output_path:
+            print(f"Moving file to {output_path}...")
+            move(str(temp_path), str(output_path))
+
+    return output_path