BTP_CODE_WITH_LLMS.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NYOJ1VTMKRoT"
      },
      "outputs": [],
      "source": [
        "# First, let's make sure we have a clean environment for FAISS\n",
        "!pip uninstall -y faiss-cpu faiss-gpu\n",
        "!pip uninstall -y faiss\n",
        "\n",
        "# Now install FAISS specifically for CPU\n",
        "!apt-get update && apt-get install -y python3-dev\n",
        "!pip install faiss-cpu --no-cache-dir\n",
        "\n",
        "# Verify FAISS installation\n",
        "import faiss\n",
        "import numpy as np\n",
        "\n",
        "# Test FAISS functionality\n",
        "print(\"\\nTesting FAISS installation:\")\n",
        "dimension = 64\n",
        "nb = 100\n",
        "xb = np.random.random((nb, dimension)).astype('float32')\n",
        "index = faiss.IndexFlatL2(dimension)\n",
        "index.add(xb)\n",
        "print(\"FAISS test successful!\")\n",
        "\n",
        "# Now let's verify we have all required components\n",
        "print(\"\\nVerifying all components:\")\n",
        "import torch\n",
        "import transformers\n",
        "import sentence_transformers\n",
        "import pandas as pd\n",
        "\n",
        "print(f\"PyTorch version: {torch.__version__}\")\n",
        "print(f\"Transformers version: {transformers.__version__}\")\n",
        "print(f\"Sentence-transformers version: {sentence_transformers.__version__}\")\n",
        "print(f\"FAISS version: {faiss.__version__}\")\n",
        "print(f\"Pandas version: {pd.__version__}\")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "import logging\n",
        "from typing import List, Dict\n",
        "import torch\n",
        "from sentence_transformers import SentenceTransformer\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "from tqdm.auto import tqdm\n",
        "import time\n",
        "\n",
        "# Setup logging\n",
        "logging.basicConfig(\n",
        "    level=logging.INFO,\n",
        "    format='%(asctime)s - %(levelname)s - %(message)s'\n",
        ")\n",
        "logger = logging.getLogger(__name__)\n",
        "\n",
        "class RobustRAG:\n",
        "    def __init__(self, data_path: str):\n",
        "        \"\"\"Initialize RAG system with component verification\"\"\"\n",
        "        self.data_path = data_path\n",
        "        self.device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "        logger.info(f\"Using device: {self.device}\")\n",
        "\n",
        "        # Initialize components as None\n",
        "        self.data = None\n",
        "        self.embedding_model = None\n",
        "        self.embeddings = None\n",
        "        self.index = None\n",
        "        self.tokenizer = None\n",
        "        self.llm_model = None\n",
        "\n",
        "        # Track component status\n",
        "        self.status = {\n",
        "            'data_loaded': False,\n",
        "            'embedding_model_loaded': False,\n",
        "            'embeddings_generated': False,\n",
        "            'faiss_initialized': False,\n",
        "            'llm_loaded': False\n",
        "        }\n",
        "\n",
        "    def verify_faiss(self):\n",
        "        \"\"\"Verify FAISS installation and functionality\"\"\"\n",
        "        try:\n",
        "            import faiss\n",
        "            # Test FAISS with small random data\n",
        "            dimension = 64\n",
        "            nb = 10\n",
        "            xb = np.random.random((nb, dimension)).astype('float32')\n",
        "            index = faiss.IndexFlatL2(dimension)\n",
        "            index.add(xb)\n",
        "            logger.info(\"FAISS verification successful\")\n",
        "            return True\n",
        "        except Exception as e:\n",
        "            logger.error(f\"FAISS verification failed: {str(e)}\")\n",
        "            return False\n",
        "\n",
        "    def load_data(self):\n",
        "        \"\"\"Load and prepare dataset with retries\"\"\"\n",
        "        max_retries = 3\n",
        "        for attempt in range(max_retries):\n",
        "            try:\n",
        "                # Read CSV file (adjust separator if needed)\n",
        "                self.data = pd.read_csv(self.data_path)  # Removed sep='\\t'\n",
        "                # Print available columns for debugging\n",
        "                logger.info(f\"Available columns: {list(self.data.columns)}\")\n",
        "\n",
        "                # Determine the correct description column\n",
        "                possible_description_columns = ['Description', 'Abstract', 'Summary']\n",
        "                description_column = None\n",
        "                for col in possible_description_columns:\n",
        "                    if col in self.data.columns:\n",
        "                        description_column = col\n",
        "                        break\n",
        "                if not description_column:\n",
        "                    raise KeyError(\"No description column found in the dataset. Checked columns: 'Description', 'Abstract', 'Summary'.\")\n",
        "\n",
        "                # Fill missing descriptions\n",
        "                self.data['Description'] = self.data[description_column].fillna('No description available')\n",
        "                logger.info(f\"Loaded dataset with {len(self.data)} entries\")\n",
        "                self.status['data_loaded'] = True\n",
        "                return True\n",
        "            except Exception as e:\n",
        "                logger.warning(f\"Data loading attempt {attempt + 1} failed: {str(e)}\")\n",
        "                # Print available columns if KeyError occurs\n",
        "                if isinstance(e, KeyError):\n",
        "                    logger.warning(f\"Available columns are: {list(self.data.columns)}\")\n",
        "                if attempt < max_retries - 1:\n",
        "                    time.sleep(1)  # Wait before retrying\n",
        "                else:\n",
        "                    logger.error(\"All data loading attempts failed\")\n",
        "                    return False\n",
        "\n",
        "    def setup_embedding_model(self):\n",
        "        \"\"\"Setup embedding model with verification\"\"\"\n",
        "        try:\n",
        "            self.embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')\n",
        "\n",
        "            # Verify model works\n",
        "            test_text = \"Test sentence for verification.\"\n",
        "            test_embedding = self.embedding_model.encode([test_text])\n",
        "            assert test_embedding.shape[1] > 0\n",
        "\n",
        "            logger.info(\"Embedding model setup successful\")\n",
        "            self.status['embedding_model_loaded'] = True\n",
        "            return True\n",
        "        except Exception as e:\n",
        "            logger.error(f\"Embedding model setup failed: {str(e)}\")\n",
        "            return False\n",
        "\n",
        "    def generate_embeddings(self):\n",
        "        \"\"\"Generate embeddings with progress tracking and error handling\"\"\"\n",
        "        if not self.status['data_loaded'] or not self.status['embedding_model_loaded']:\n",
        "            logger.error(\"Cannot generate embeddings: prerequisites not met\")\n",
        "            return False\n",
        "\n",
        "        try:\n",
        "            descriptions = self.data['Description'].tolist()\n",
        "            self.embeddings = []\n",
        "            batch_size = 32\n",
        "\n",
        "            for i in tqdm(range(0, len(descriptions), batch_size), desc=\"Generating embeddings\"):\n",
        "                batch = descriptions[i:i + batch_size]\n",
        "                try:\n",
        "                    batch_embeddings = self.embedding_model.encode(\n",
        "                        batch,\n",
        "                        convert_to_numpy=True,\n",
        "                        show_progress_bar=False\n",
        "                    )\n",
        "                    self.embeddings.append(batch_embeddings)\n",
        "                except Exception as e:\n",
        "                    logger.warning(f\"Failed to process batch {i//batch_size}: {str(e)}\")\n",
        "                    # Create zero embeddings for failed batch\n",
        "                    batch_embeddings = np.zeros((len(batch), self.embedding_model.get_sentence_embedding_dimension()))\n",
        "                    self.embeddings.append(batch_embeddings)\n",
        "\n",
        "            self.embeddings = np.vstack(self.embeddings)\n",
        "            self.status['embeddings_generated'] = True\n",
        "            logger.info(f\"Generated embeddings of shape {self.embeddings.shape}\")\n",
        "            return True\n",
        "        except Exception as e:\n",
        "            logger.error(f\"Embedding generation failed: {str(e)}\")\n",
        "            return False\n",
        "\n",
        "    def setup_faiss(self):\n",
        "        \"\"\"Setup FAISS index with verification\"\"\"\n",
        "        if not self.status['embeddings_generated']:\n",
        "            logger.error(\"Cannot setup FAISS: embeddings not generated\")\n",
        "            return False\n",
        "\n",
        "        try:\n",
        "            import faiss\n",
        "            embedding_dim = self.embeddings.shape[1]\n",
        "            self.index = faiss.IndexFlatL2(embedding_dim)\n",
        "            self.index.add(self.embeddings.astype('float32'))\n",
        "\n",
        "            # Verify index\n",
        "            test_query = self.embeddings[0:1]\n",
        "            distances, indices = self.index.search(test_query, 1)\n",
        "            assert indices.shape == (1, 1)\n",
        "\n",
        "            self.status['faiss_initialized'] = True\n",
        "            logger.info(\"FAISS index created and verified\")\n",
        "            return True\n",
        "        except Exception as e:\n",
        "            logger.error(f\"FAISS setup failed: {str(e)}\")\n",
        "            return False\n",
        "\n",
        "    def setup_llm(self):\n",
        "        \"\"\"Setup LLM with verification\"\"\"\n",
        "        try:\n",
        "            self.tokenizer = AutoTokenizer.from_pretrained(\"meta-llama/Llama-2-7b-hf\")\n",
        "            self.llm_model = AutoModelForCausalLM.from_pretrained(\n",
        "                \"meta-llama/Llama-2-7b-hf\",\n",
        "                device_map=\"auto\",\n",
        "                torch_dtype=torch.float16,\n",
        "                low_cpu_mem_usage=True\n",
        "            )\n",
        "\n",
        "            # Verify tokenizer and model\n",
        "            test_input = \"Test input.\"\n",
        "            tokens = self.tokenizer(test_input, return_tensors=\"pt\")\n",
        "            assert tokens is not None\n",
        "\n",
        "            self.status['llm_loaded'] = True\n",
        "            logger.info(\"LLM setup successful\")\n",
        "            return True\n",
        "        except Exception as e:\n",
        "            logger.error(f\"LLM setup failed: {str(e)}\")\n",
        "            return False\n",
        "\n",
        "    def setup(self):\n",
        "        \"\"\"Complete setup with component verification\"\"\"\n",
        "        steps = [\n",
        "            (self.verify_faiss, \"FAISS verification\"),\n",
        "            (self.load_data, \"Data loading\"),\n",
        "            (self.setup_embedding_model, \"Embedding model setup\"),\n",
        "            (self.generate_embeddings, \"Embedding generation\"),\n",
        "            (self.setup_faiss, \"FAISS setup\"),\n",
        "            (self.setup_llm, \"LLM setup\")\n",
        "        ]\n",
        "\n",
        "        success = True\n",
        "        for step_func, step_name in steps:\n",
        "            logger.info(f\"Starting {step_name}...\")\n",
        "            if not step_func():\n",
        "                success = False\n",
        "                logger.warning(f\"{step_name} failed or partially completed\")\n",
        "            else:\n",
        "                logger.info(f\"{step_name} completed successfully\")\n",
        "\n",
        "        return success\n",
        "\n",
        "    def query(self, text: str, top_k: int = 3):\n",
        "        \"\"\"Query system with component checks\"\"\"\n",
        "        if not all([self.status['embeddings_generated'], self.status['faiss_initialized']]):\n",
        "            return {\"error\": \"System not fully initialized\", \"results\": []}\n",
        "\n",
        "        try:\n",
        "            query_embedding = self.embedding_model.encode([text])\n",
        "            distances, indices = self.index.search(query_embedding.astype('float32'), top_k)\n",
        "\n",
        "            results = []\n",
        "            for idx, distance in zip(indices[0], distances[0]):\n",
        "                results.append({\n",
        "                    'text': self.data['Description'].iloc[idx],\n",
        "                    'title': self.data['Title'].iloc[idx] if 'Title' in self.data.columns else 'No title available',\n",
        "                    'accession_id': self.data['Accession ID'].iloc[idx] if 'Accession ID' in self.data.columns else 'No accession ID available',\n",
        "                    'distance': float(distance)\n",
        "                })\n",
        "\n",
        "            return {\"error\": None, \"results\": results}\n",
        "        except Exception as e:\n",
        "            logger.error(f\"Query failed: {str(e)}\")\n",
        "            return {\"error\": str(e), \"results\": []}\n",
        "\n",
        "def main():\n",
        "    # Initialize system\n",
        "    rag = RobustRAG('/content/bioproject_Dengue_Human_sequencing.csv')\n",
        "\n",
        "    # Setup system\n",
        "    if rag.setup():\n",
        "        logger.info(\"System setup complete!\")\n",
        "\n",
        "        # Test query PRJNA1063364;\n",
        "        query = \"RNA-seq and ATAC-seq analysis of iPSC macrophages\"\n",
        "        #query =\"Macrophage differentiation and polarization using iPSCs\"\n",
        "        #query =\"Flow cytometry and transcriptomics in macrophages\"\n",
        "        #query =\"Chromatin accessibility profiling in virus-infected macrophages\"\n",
        "        #query = \"RNA secondary structure and gene expression in single cells.\"\n",
        "        results = rag.query(query)\n",
        "\n",
        "        if results[\"error\"] is None:\n",
        "            print(\"\\nQuery results:\")\n",
        "            for result in results[\"results\"]:\n",
        "                print(f\"\\nAccession ID: {result['accession_id']}\")\n",
        "                print(f\"Title: {result['title']}\")\n",
        "                print(f\"Distance: {result['distance']:.4f}\")\n",
        "                print(f\"Text: {result['text'][:200]}...\")\n",
        "        else:\n",
        "            print(f\"Query failed: {results['error']}\")\n",
        "    else:\n",
        "        logger.error(\"System setup incomplete. Check logs for details.\")\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main()\n"
      ],
      "metadata": {
        "id": "C1PQnvC-KVxw"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}