tutorials/generate_and_store_embeddings_genome_huggingface.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e1d8f8f8",
   "metadata": {},
   "source": [
    "# Kikiola Genome Embedding with Hugging Face 🤗"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7aa7775b",
   "metadata": {},
   "source": [
    "## Installing Required Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7f9c6d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install transformers requests torch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3b0c8e6",
   "metadata": {},
   "source": [
    "## Starting the Server"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1234abcd",
   "metadata": {},
   "source": [
    "Before running the Kikiola Genome Embedding code, make sure to start the server by running the following command in your terminal. This command will start the server that will handle the storage of the generated embeddings.\n",
    "\n",
    "```sh\n",
    "go run cmd/main.go\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3b0c8e6",
   "metadata": {},
   "source": [
    "## Kikiola Genome Embedding Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5f8f8f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "import sys\n",
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModel\n",
    "\n",
    "class KikiolaGenomeEmbedding:\n",
    "    def __init__(self):\n",
    "        self.tokenizer = AutoTokenizer.from_pretrained(\"sentence-transformers/all-MiniLM-L6-v2\")\n",
    "        self.model = AutoModel.from_pretrained(\"sentence-transformers/all-MiniLM-L6-v2\")\n",
    "        self.document_text = \"\"\n",
    "        self.embeddings = []\n",
    "\n",
    "    def load_genome_sequences(self, gene_id):\n",
    "        server = \"http://rest.ensembl.org\"\n",
    "        ext = f\"/genetree/member/id/{gene_id}?\"\n",
    "        r = requests.get(server + ext, headers={\"Content-Type\": \"application/json\"})\n",
    "        if not r.ok:\n",
    "            r.raise_for_status()\n",
    "            sys.exit()\n",
    "        decoded = r.json()\n",
    "        self.extract_genome_sequences(decoded)\n",
    "\n",
    "    def extract_genome_sequences(self, data):\n",
    "        if isinstance(data, dict):\n",
    "            if 'sequence' in data and 'mol_seq' in data['sequence']:\n",
    "                genome_sequence = data['sequence']['mol_seq']['seq']\n",
    "                self.document_text += genome_sequence + '\\n'\n",
    "            else:\n",
    "                for value in data.values():\n",
    "                    self.extract_genome_sequences(value)\n",
    "        elif isinstance(data, list):\n",
    "            for item in data:\n",
    "                self.extract_genome_sequences(item)\n",
    "\n",
    "    def generate_embeddings(self):\n",
    "        chunk_size = 8000\n",
    "        chunks = [self.document_text[i:i + chunk_size] for i in range(0, len(self.document_text), chunk_size)]\n",
    "        self.embeddings = []\n",
    "        for chunk in chunks:\n",
    "            encoded_input = self.tokenizer(chunk, return_tensors='pt', truncation=True, padding=True)\n",
    "            with torch.no_grad():\n",
    "                model_output = self.model(**encoded_input)\n",
    "            embeddings = model_output.last_hidden_state.mean(dim=1).tolist()\n",
    "            self.embeddings.extend(embeddings)\n",
    "\n",
    "    def store_embeddings(self):\n",
    "        for i, embedding in enumerate(self.embeddings):\n",
    "            vector_data = {\n",
    "                \"id\": f\"genome_sequence_{i}\",\n",
    "                \"embedding\": embedding,\n",
    "                \"metadata\": {\n",
    "                    \"name\": f\"Genome Sequence Embeddings - Part {i+1}\",\n",
    "                    \"category\": \"genome_sequence\"\n",
    "                }\n",
    "            }\n",
    "            print(f\"Vector data for Part {i+1}: {vector_data}\")\n",
    "            response = requests.post(\"http://localhost:3400/vectors\", json=vector_data)\n",
    "            if response.status_code == 200:\n",
    "                print(f\"Embeddings stored for Part {i+1}. Status code: {response.status_code}\")\n",
    "            else:\n",
    "                print(f\"Error storing embeddings for Part {i+1}. Status code: {response.status_code}\")\n",
    "                print(f\"Error response: {response.text}\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    gene_id = \"ENSG00000157764\"\n",
    "    embeddings_generator = KikiolaGenomeEmbedding()\n",
    "    embeddings_generator.load_genome_sequences(gene_id)\n",
    "    embeddings_generator.generate_embeddings()\n",
    "    embeddings_generator.store_embeddings()\n",
    "    print(\"Kikiola Embeddings Completed.\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}