pmc-retrieval-demo.ipynb

{"cells":[{"cell_type":"markdown","id":"864ffbd4-ebb2-42c9-997f-0ef68f4d0fb8","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}}},"source":["# PMC RAG Demo and Deep Dive\n","## Why RAG\n","By embracing RAG, you can unlock a range of benefits for your organization:\n","\n"," * Improved decision-making: Accessing rights and trustworthy information empowers better choices and strategies.\n"," * Enhanced customer experience: Delivering reliable answers and insights builds trust and satisfaction.\n"," * Reduced risk and compliance: Curated data sources minimize the risk of misinformation and ensure compliance with regulations.\n"," * Increased efficiency: Streamlining access to information saves time and resources.\n"," \n","The beauty of RAG lies in its focus on data quality, not just data quantity. We're moving beyond the “bigger is better” mentality of massive models trained on internet data that often include misinformation and biases. RAG puts the emphasis on smaller, more valuable models that use curated, trustworthy data sources.\n"]},{"cell_type":"markdown","id":"f983c8f8-d56d-4f5b-a51e-e5df2bb33982","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}}},"source":["## The Standard Rag Workflow Empowerd by TileDB\n","\n","1. **User:** Uploads documents, and the system converts them into vectors (numeric representations) using sentence embeddings. In our case, the user submits a text file that the system then uses to scrape Pub Med Central.\n","2. **User:** Stores these document vectors, along with the documents and metadata, into TileDB (a smart database for storing vectors).\n","3. **User:** Asks a question.\n","4. **Embedding Model:** Processes the user's question by embedding it into a vector and sends this vector to TileDB.\n","5. **VectorDB:** Searches through the stored document vectors and retrieves the most relevant documents.\n","6. **Retriever:** Takes these relevant documents and constructs a new query for the LLM, instructing it to use these documents as context.\n","7. **LLM:** Uses the relevant documents as context to generate and deliver the final answer back to the user."]},{"cell_type":"markdown","id":"00897ed6-52ec-4d4a-b116-2a286b990f68","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}}},"source":["## PHASE 1 The Notebook and Ingestion Pipelines\n","Below is code that you can run in order. The ingestion process will take a bit, but we will discuss further in this document how you can improve and iterate on this example depending on your objectives."]},{"cell_type":"markdown","id":"f4899111-8920-4850-99f8-666d406e7154","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}}},"source":["### Imports\n","The below imports are for our local file push to our remote repo as well as the `initialize_step()` that will build our end to end pipeline for us and submit a dag directly. "]},{"cell_type":"code","execution_count":null,"id":"c67374d0-369e-44e1-8a61-036b9a20f7b3","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}},"trusted":true},"outputs":[],"source":["import requests\n","import os\n","import pandas as pd\n","import tarfile\n","import urllib.request\n","import xml.etree.ElementTree as ET\n","import hashlib\n","import subprocess\n","import shutil\n","import math\n","import warnings\n","import time\n","from datetime import datetime\n","from typing import Optional, List, Union, Dict, Iterator\n","\n","# Langchain Imports\n","from langchain._api import LangChainDeprecationWarning\n","warnings.filterwarnings(\"ignore\")\n","warnings.filterwarnings(\"ignore\", category=DeprecationWarning) \n","warnings.filterwarnings(\"ignore\", category=LangChainDeprecationWarning)\n","from langchain.document_loaders import PyPDFLoader\n","from langchain.document_loaders.parsers.pdf import PyPDFParser\n","from langchain.embeddings import HuggingFaceEmbeddings\n","from langchain.embeddings.openai import OpenAIEmbeddings\n","from langchain.document_loaders.pdf import BasePDFLoader\n","from langchain.document_loaders.blob_loaders import Blob\n","from langchain.schema import Document\n","from langchain.memory import ConversationSummaryMemory\n","from langchain.chains import ConversationalRetrievalChain, ConversationChain\n","from langchain.vectorstores.tiledb import TileDB\n","\n","# TileDB Imports \n","import tiledb\n","import tiledb.cloud\n","from tiledb.cloud.dag import dag\n","from tiledb.vector_search.object_api import object_index\n","from tiledb.vector_search.object_readers import DirectoryTextReader\n","from tiledb.vector_search.embeddings import SentenceTransformersEmbedding, LangChainEmbedding\n","from tiledb.cloud import groups\n","from langchain_community.embeddings import HuggingFaceEmbeddings\n","from sentence_transformers import SentenceTransformer  # Ensure this is correctly imported\n","# Environment Variables \n","os.environ['TOKENIZERS_PARALLELISM'] = \"true\""]},{"cell_type":"markdown","id":"a39d443d-6d54-4d4a-9107-b5a8ce69ad81","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}}},"source":["### Credentials\n","The below cell is so we can cache our credentials during an initial push for our steps. After you push the local file, you may need to manually enter creds and push the file from the terminal. Afterward, the credentials should be cached and you can run without issues. "]},{"cell_type":"code","execution_count":null,"id":"ab801d52-fc16-4a07-ad02-4c75b8e301dd","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}},"trusted":true},"outputs":[],"source":["#before running below please run this. The first push to the repo (if necessary) may fail and you will need to manually push thie file.\n","#after, the credentials should be temporarily cached. \n","#!git config --global user.email \"you@example.com\"\n","#!git config --global user.name \"Your Name\"\n","#!git config --global credential.helper cache\n","#!git config --global credential.helper 'cache --timeout=3600'\n","#!git config --global credential.helper store"]},{"cell_type":"markdown","id":"558a5935-9d69-44f8-96cb-a9ad27963199","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}}},"source":["### Upload \"pipeline\" Helpers\n","For now, this is a quick simulation of a pipeline for updating the local file. The first cell is full ofhelper functions for our notebook pipeline (ran in our notebook). "]},{"cell_type":"code","execution_count":null,"id":"12c33aa8-e715-46ba-9495-706a63e1194c","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}},"trusted":true},"outputs":[],"source":["def add_and_commit_files(message: str):\n","    \"\"\"\n","    Adds all files to the Git staging area, commits them with a provided message, and pushes the changes to the remote repository.\n","    \n","    :param message: Commit message to be used in the Git commit.\n","    \"\"\"\n","    \n","    # Print the current working directory\n","    print(f\"Working directory is {os.getcwd()}\")\n","\n","    # Stage all changes (new, modified, deleted) in the current Git repository\n","    subprocess.run([\"git\", \"add\", \"-A\"])\n","\n","    # Commit the staged changes with the provided commit message\n","    subprocess.run([\"git\", \"commit\", \"-m\", f\"{message}\"])\n","\n","    # Push the committed changes to the remote repository (origin/master by default)\n","    subprocess.run([\"git\", \"push\"])\n","\n","def hash_file(file_path: str) -> str:\n","    \"\"\"\n","    Computes the SHA256 hash of the contents of a file.\n","\n","    :param file_path: Path to the file to be hashed.\n","    :return: The SHA256 hash of the file contents as a hex string.\n","    \"\"\"\n","    \n","    # Create a new SHA256 hash object\n","    hasher = hashlib.sha256()\n","    \n","    # Open the file in binary read mode\n","    with open(file_path, 'rb') as f:\n","        # Read the file contents and update the hash object\n","        buffer = f.read()\n","        hasher.update(buffer)\n","    \n","    # Return the hexadecimal digest of the hash\n","    return hasher.hexdigest()\n","\n","def estimate_resources(job_df: pd.DataFrame):\n","    \"\"\"\n","    Estimates the CPU and memory usage required for processing a given DataFrame, with an additional 2 GB overhead, \n","    and rounds the total memory usage to the nearest GB.\n","\n","    :param job_df: A pandas DataFrame representing job data.\n","    :return: A tuple containing the estimated number of CPUs per job and the total memory usage in GB (rounded).\n","    \"\"\"\n","    \n","    # Print the data types of each column in the DataFrame for reference\n","    print(\"Data types of the DataFrame:\")\n","    print(job_df.dtypes)\n","\n","    # Calculate the memory usage of each column in the DataFrame in MB (deep=True considers the actual memory usage)\n","    memory_usage_per_column = job_df.memory_usage(deep=True) / (1024 ** 2)  # Convert bytes to MB\n","    total_memory_usage = memory_usage_per_column.sum()  # Sum of the memory usage of all columns\n","    \n","    # Convert total memory usage to GB and add an additional 2 GB overhead for processing\n","    total_memory_usage_gb = total_memory_usage / 1024  # Convert MB to GB\n","    total_memory_usage_gb += 2  # Add 2 GB overhead for processing\n","    \n","    # Round up the total memory usage to the nearest GB\n","    rounded_memory_usage_gb = math.ceil(total_memory_usage_gb)\n","    \n","    # Print the memory usage for each column and the total estimated memory usage\n","    print(\"Memory usage per column (MB):\")\n","    print(memory_usage_per_column)\n","    \n","    print(f\"Total memory usage for job (GB, rounded to the nearest GB, with overhead): {rounded_memory_usage_gb}\")\n","\n","    # Estimate the CPU usage per job (adjustable based on job complexity)\n","    cpu_per_job = 1  # Example: assuming 1 CPU per job\n","    \n","    # Return the estimated CPU count and the rounded memory usage in GB\n","    return cpu_per_job, rounded_memory_usage_gb\n"]},{"cell_type":"markdown","id":"848ad50b-5fc7-4c57-9b42-038dc4f44807","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}}},"source":["### The \"Upload Pipeline\" \n","This next cell is similar to a DevOps runner pipeline where a user would commit an updated file and the run would create a container with a run hash as the tag. The goal here is to determine if there is a change in the local file and update the hash. The next stage of the pipeline will/would use the hash to determine if a run is necessary. The frequency of checking a run really depends on the frequency of the file update and how we want to tune it to adjust the quality of our LLMs outputs. Fututre state this could be a webhook upon git update, or s3 storage. The benefit of git for this is the \"GitOps\" like workflow of tracking changes to our ingestion documents file for RAG and then using that knowledge to understand the impact on our model outputs. \n","\n","#### The Code Below in Plain English\n","\n","1. **Change to Home Directory:** It starts by navigating to the user's home directory.\n","\n","2. **Clone or Pull Repository:** It checks if a given repository (based on the URL) already exists locally. If it exists, it updates the repository by pulling the latest changes. If it doesn't exist, it clones the repository from the given URL.\n","\n","3. **Hash a Local File:** It calculates a hash (unique identifier) of a local file's contents to check if it has been modified.\n","\n","4. **Compare the Hash:** It checks whether a previously saved hash (in a hash.txt file) exists. If it does, the new hash is compared to the saved one to determine if the file has changed.\n","\n","5. **Skip or Proceed:** If the file hasn't changed (i.e., the current hash matches the previous one), it skips any further action. If the file has changed, it proceeds.\n","\n","6. **Update the Hash File:** It writes the new hash into a hash.txt file in the repository.\n","\n","7. **Commit Changes to Git:** Finally, it navigates to the repository directory, adds, commits, and pushes the new or modified file and updated hash to the Git repository"]},{"cell_type":"code","execution_count":null,"id":"4578b3e2-d8ff-4a0b-9e9a-7909dd0a0997","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}},"trusted":true},"outputs":[],"source":["def handle_local_file(repo_url: str, local_file_name: str, hash_file_name: str):\n","    \"\"\"\n","    Clones or pulls a Git repository, checks if a local file has changed by comparing its hash with a stored hash,\n","    and if the file has been modified, pushes the updated file and hash to the repository.\n","    \n","    :param repo_url: The URL of the Git repository.\n","    :param local_file_name: The local file whose hash will be checked for modifications.\n","    :param hash_file_name: The file in the repo that stores the previous hash of the local file.\n","    \"\"\"\n","    \n","    # Step 1: Navigate to the home directory to clone or pull the repo\n","    home_directory = os.path.expanduser(\"~\")  # Get the user's home directory\n","    os.chdir(home_directory)  # Change the working directory to the home directory\n","    \n","    # Extract the repository name from the repo URL (assumes .git format at the end)\n","    repo_name = repo_url.split('/')[-1].replace('.git', '')\n","    \n","    # Check if the repository already exists locally\n","    if os.path.exists(repo_name): \n","        # If the repo exists, navigate into it and pull the latest changes\n","        os.chdir(repo_name)\n","        subprocess.run([\"git\", \"pull\"])  # Pull latest changes from the remote repo\n","        os.chdir(\"..\")  # Go back to the home directory after pulling changes\n","    else:\n","        # If the repo doesn't exist, clone it from the provided URL\n","        subprocess.run([\"git\", \"clone\", repo_url])  # Clone the repository\n","    \n","    # Get the full path to the local repository\n","    local_repo_path = os.path.join(os.getcwd(), repo_name)\n","    print(f\"Created {local_repo_path}\")  # Print the path to the repo\n","    \n","    # Step 2: Hash the contents of the local file\n","    current_hash = hash_file(local_file_name)  # Hash the local file\n","    \n","    # Step 3: Check if the hash file exists in the repository and compare hashes\n","    hash_file_path = os.path.join(local_repo_path, hash_file_name)  # Path to the hash file in the repo\n","    \n","    if os.path.exists(hash_file_path):\n","        # If the hash file exists, read the previous hash\n","        with open(hash_file_path, 'r') as f:\n","            previous_hash = f.read().strip()  # Strip any extra whitespace\n","    else:\n","        previous_hash = \"\"  # If the hash file doesn't exist, assume no previous hash\n","    \n","    # Step 4: Compare the current hash with the previous hash\n","    if current_hash == previous_hash:\n","        # If the hashes match, the file hasn't changed\n","        print(\"File has not changed, skipping submission.\")\n","        return\n","    else:\n","        # If the hashes differ, the file has changed, so proceed with updating the repo\n","        print(\"File has changed, proceeding with submission.\")\n","    \n","    # Step 5: Update the hash file in the repository with the new hash\n","    with open(hash_file_path, 'w') as f:\n","        f.write(current_hash)  # Write the new hash to the file\n","    \n","    # Step 6: Commit the changes (the updated file and the new hash) to the Git repository\n","    os.chdir(local_repo_path)  # Change directory to the repo\n","    add_and_commit_files(\"added local articles file to the repository\")  # Stage, commit, and push the changes\n"]},{"cell_type":"code","execution_count":null,"id":"46ece7fb-6197-40c6-be21-31cc2c30b5fd","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}},"trusted":true},"outputs":[],"source":["# Example usage\n","#handle_local_file(\n","#    repo_url=\"https://github.com/TileDB-Inc/pmc-llm.git\", \n","#    local_file_name=\"rag-article-list.txt\", \n","#    hash_file_name=\"hash.txt\"\n","#)"]},{"cell_type":"markdown","id":"aa8d563e-8148-4ac0-8ab7-d94ed0e713c9","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}}},"source":["### Ingestion Tasks for the TileDB DAG"]},{"cell_type":"code","execution_count":null,"id":"2e31c91d-8683-4217-b65f-8e6b22eeff0f","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}},"trusted":true},"outputs":[],"source":["def consolidate_chunks(total_jobs, bucket_region, object_path, prefix):\n","    \"\"\"\n","    Consolidates a list of chunked files into a single file in S3 and deletes the chunked files afterward.\n","\n","    :param total_jobs: Number of chunked files to process\n","    :param bucket_region: The S3 bucket region\n","    :param object_path: The S3 path to the files\n","    :param prefix: The prefix of the chunked files\n","    \"\"\"\n","    import tiledb\n","    \n","    # Create a TileDB context with the specified S3 region\n","    ctx = tiledb.Ctx({\"vfs.s3.region\": bucket_region})\n","    # Initialize the TileDB VFS (Virtual File System) to interact with S3\n","    vfs = tiledb.VFS(ctx=ctx)\n","    \n","    # Initialize tmp_data as an empty bytes object to store the consolidated data\n","    tmp_data = b\"\"  # Use bytes since the files are being read in binary mode\n","    \n","    # Loop through all the chunked files and concatenate them\n","    for i in range(total_jobs):\n","        # Construct the file path for each chunk\n","        file_path = f\"{object_path}/history/{prefix}_{i}\"\n","        \n","        # Check if the file exists in the S3 bucket\n","        if vfs.is_file(file_path):\n","            # If the file exists, open it in binary read mode and append its contents\n","            with vfs.open(file_path, 'rb') as f:\n","                tmp_data += f.read()  # Read and concatenate the binary content\n","            \n","            # Delete the chunked file after reading\n","            vfs.remove_file(file_path)\n","            print(f\"Deleted chunk file: {file_path}\")\n","        else:\n","            # If the file doesn't exist, print a message and continue to the next file\n","            print(f\"No such file path {file_path}. Moving onto the next file.\")\n","            continue\n","    \n","    # Now `tmp_data` contains the concatenated data from all chunked files\n","    # Define the path for the consolidated file in the S3 bucket\n","    consolidated_file_path = f\"{object_path}/history/consolidated_{prefix}\"\n","    \n","    # Open the consolidated file in binary write mode and write the consolidated data\n","    with vfs.open(consolidated_file_path, 'wb') as f:\n","        f.write(tmp_data)  # Write the concatenated binary data to the new file\n","    \n","    # Print a success message with the path of the consolidated file\n","    print(f\"Consolidation complete. Consolidated file uploaded to: {consolidated_file_path}\")\n"]},{"cell_type":"code","execution_count":null,"id":"19dbc31b-b3fa-44bd-a35f-666022bbe261","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}},"trusted":true},"outputs":[],"source":["def delete_unwanted_files(bucket_region, bucket_path):\n","    \"\"\"\n","    Deletes files in an S3 bucket that are not listed in the 'consolidated_file_history' file.\n","    \n","    :param bucket_region: The S3 bucket region.\n","    :param bucket_path: The S3 path to the bucket.\n","    \"\"\"\n","    import tiledb\n","    import os\n","    \n","    # Initialize TileDB context and VFS with the specified S3 region\n","    ctx = tiledb.Ctx({\"vfs.s3.region\": bucket_region})\n","    vfs = tiledb.VFS(ctx=ctx)\n","    \n","    # Path to the 'consolidated_file_history' file in S3\n","    consolidated_history_path = os.path.join(bucket_path, \"history/consolidated_file_history\")\n","    \n","    # Step 1: Load the list of valid files from 'consolidated_file_history'\n","    valid_files = set()  # Initialize an empty set to store valid file names\n","    if vfs.is_file(consolidated_history_path):\n","        # If 'consolidated_file_history' exists, open and read it\n","        with vfs.open(consolidated_history_path, 'rb') as f:\n","            # Read each line, decode from bytes to string, and strip newline characters\n","            valid_files = set(line.decode('utf-8').strip() for line in f.read().splitlines())\n","    else:\n","        # Raise an error if the 'consolidated_file_history' file is not found\n","        raise FileNotFoundError(f\"{consolidated_history_path} not found.\")\n","    \n","    # Step 2: List all files in the S3 bucket\n","    all_files = []  # Initialize a list to store all file paths\n","    if vfs.is_dir(bucket_path):\n","        # If the bucket is a directory, list all files\n","        for file in vfs.ls(bucket_path):\n","            if not file.endswith('/'):  # Skip directories\n","                all_files.append(file)\n","    \n","    print(\"Begin the search!\")\n","    \n","    # Step 3: Delete unwanted files\n","    for file_path in all_files:\n","        # Extract the file name from the full path\n","        file_name = os.path.basename(file_path)\n","        \n","        # Skip the 'consolidated_file_history' file to avoid deleting it\n","        if file_name == \"consolidated_file_history\":\n","            print(f\"Skipping deletion of {file_name}, as it is the consolidated file.\")\n","            continue  # Skip further checks for this file\n","        \n","        # Check if the file name is NOT in the consolidated list, delete if not\n","        if file_name not in valid_files:\n","            print(f\"Deleting {file_path} because it's not in consolidated_file_history.\")\n","            vfs.remove_file(file_path)  # Remove the file from S3\n","        else:\n","            print(f\"Keeping {file_path}, it's in consolidated_file_history.\")\n"]},{"cell_type":"code","execution_count":null,"id":"c0954fae-8a82-42cb-81b0-c55a67c619a7","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}},"tags":[],"trusted":true},"outputs":[],"source":["def pmid_ingestion(job_df, object_directory: str, bucket_region: str, job_id: int, index_uri: str, embedding_model_name: str):\n","    \"\"\"\n","    Main function to handle the ingestion of articles using PubMed API and TileDB VFS.\n","    \n","    :param job_df: DataFrame containing PubMed PMIDs and gene-disease data.\n","    :param object_directory: Path to the S3 bucket for file storage.\n","    :param bucket_region: Region of the S3 bucket.\n","    :param job_id: ID of the job to track file history.\n","    :param index_uri: TileDB index URI where the articles will be stored.\n","    :param embedding_model_name: Name of the embedding model to store as metadata in TileDB.\n","    \"\"\"\n","    import logging\n","    import time\n","    from datetime import datetime\n","    import os\n","    import urllib.request\n","    import xml.etree.ElementTree as ET\n","    import requests\n","    import tiledb\n","    from langchain_community.vectorstores import TileDB\n","    from langchain_community.embeddings import HuggingFaceEmbeddings\n","    from langchain.schema import Document\n","    from typing import List, Iterator, Optional, Union\n","    import tarfile\n","    import glob\n","    from typing import Optional, Union, Dict, List, Iterator\n","    import tiledb\n","\n","    \n","\n","    now = datetime.now()\n","    ctx = tiledb.Ctx({\"vfs.s3.region\": bucket_region})\n","    vfs = tiledb.VFS(ctx=ctx)\n","\n","    try:\n","        embeddings = HuggingFaceEmbeddings(model_name=embedding_model_name)\n","    except Exception as e:\n","        raise RuntimeError(f\"Error initializing HuggingFace embeddings with model {embedding_model_name}: {e}\")\n","\n","    # Initialize the file history log for the job\n","    new_file_path = f\"{object_directory}/history/file_history_{job_id}\"\n","    missed_file_path = f\"{object_directory}/history/missed_paper_history_{job_id}\"\n","\n","    with vfs.open(new_file_path, 'wb') as f:\n","        log_entry = f\"Run launched at {now.strftime('%Y-%m-%d %H:%M:%S')}\\n\"\n","        f.write(log_entry.encode('utf-8'))\n","        \n","    with vfs.open(missed_file_path, 'wb') as f:\n","        log_entry = f\"Run launched at {now.strftime('%Y-%m-%d %H:%M:%S')}\\n\"\n","        f.write(log_entry.encode('utf-8'))        \n","\n","    start_time = time.time()\n","\n","    def convert_pmid_to_pmcid_and_ingest(pmid):\n","        \"\"\"Convert PMID to PMCID and attempt to download the article or abstract.\"\"\"\n","        url = \"https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi\"\n","        params = {\"db\": \"pubmed\", \"id\": pmid, \"retmode\": \"json\", \"tool\": \"your_tool_name\", \"email\": \"your_email@example.com\"}\n","        max_retries = 5\n","        retry_delay = 5\n","\n","        for attempt in range(max_retries):\n","            try:\n","                response = requests.get(url, params=params)\n","                if response.status_code == 200:\n","                    data = response.json()\n","                    pmid_str = str(pmid)\n","\n","                    # Check if the result contains the expected data\n","                    if \"result\" in data and pmid_str in data[\"result\"]:\n","                        result = data[\"result\"][pmid_str]\n","                        pmcid = None\n","\n","                        # Check if a PMCID is available in the article IDs\n","                        for article_id in result.get('articleids', []):\n","                            if article_id['idtype'] == 'pmc':\n","                                pmcid = article_id['value']\n","                                if download_pmc_article(pmcid):\n","                                    add_to_new_file(f\"{pmcid}.pdf\", \"file_history\")\n","                                    return  # Exit the function once successfully processed\n","                                else:\n","                                    add_to_new_file(f\"{pmcid}.pdf\", \"missed_paper_history\")\n","                                    return  # Exit after processing (even if unsuccessful)\n","\n","                        # If no PMCID, try fetching the abstract\n","                        if fetch_abstract(pmid):\n","                            add_to_new_file(f\"{pmid}_abstract.txt\", \"file_history\")\n","                            return  # Exit after successfully fetching the abstract\n","                        else:\n","                            add_to_new_file(f\"{pmid}_abstract.txt\", \"missed_paper_history\")\n","                            return  # Exit after failing to fetch the abstract\n","\n","                    # If no valid result is found in the response, break the loop and exit\n","                    print(f\"No valid result found for PMID {pmid}. Exiting.\")\n","                    return\n","\n","                elif response.status_code == 429:\n","                    print(f\"Rate limit hit for PMID {pmid}, retrying in {retry_delay} seconds...\")\n","                    time.sleep(retry_delay)\n","                    retry_delay *= 2\n","\n","            except requests.exceptions.ConnectionError as e:\n","                print(f\"Connection error for PMID {pmid} on attempt {attempt + 1}. Retrying...\")\n","                time.sleep(retry_delay)\n","                retry_delay *= 2\n","\n","        print(f\"Failed to retrieve data for PMID {pmid} after {max_retries} attempts.\")\n","\n","    def download_pmc_article(pmcid):\n","        \"\"\"\n","        Download the article package by PMCID and upload the PDF using TileDB VFS.\n","        \"\"\"\n","        file_name = os.path.join(object_directory, f\"{pmcid}.pdf\")\n","        \n","        # Attempt ingestion regardless of whether the file exists\n","        if vfs.is_file(file_name):\n","            print(f\"File{pmcid}.pdf already exists.\")\n","            return True\n","        else:\n","            print(f\"File {pmcid}.pdf does not exist, downloading...\")\n","\n","            url = f\"https://www.ncbi.nlm.nih.gov/pmc/utils/oa/oa.fcgi?id={pmcid}\"\n","            max_retries = 5\n","            retry_delay = 5\n","\n","            for attempt in range(max_retries):\n","                try:\n","                    response = urllib.request.urlopen(url)\n","                    response_content = response.read()\n","\n","                    if response.getcode() == 200:\n","                        root = ET.fromstring(response_content)\n","                        error_element = root.find('.//error')\n","\n","                        if error_element is not None:\n","                            return False\n","                        else:\n","                            # Download and extract article package\n","                            return download_and_extract_article(root, pmcid)\n","                    elif response.status_code == 429:\n","                        time.sleep(retry_delay)\n","                        retry_delay *= 2\n","\n","                except requests.exceptions.ConnectionError as e:\n","                    time.sleep(retry_delay)\n","                    retry_delay *= 2\n","\n","        return True\n","    \n","    def download_and_extract_article(root, pmcid):\n","        \"\"\"\n","        Download and extract the article tarball.\n","        \n","        :param root: XML root element.\n","        :param pmcid: PubMed Central ID.\n","        \"\"\"\n","        import shutil\n","        records = root.find('records')\n","        if records is not None:\n","            record = records.find(f'record[@id=\"{pmcid}\"]')\n","            link = record.find('link[@format=\"tgz\"]')\n","            if link is not None:\n","                tar_url = link.get('href')\n","                tar_file_name = f\"{pmcid}.tar.gz\"\n","\n","                print(f\"Downloading {tar_url}...\")\n","                urllib.request.urlretrieve(tar_url, tar_file_name)\n","\n","                print(f\"Extracting {tar_file_name}...\")\n","                with tarfile.open(tar_file_name, 'r:gz') as tar:\n","                    tar.extractall(pmcid)\n","                os.remove(tar_file_name)\n","\n","                os.chdir(pmcid)\n","                pdf_files = glob.glob(\"**/*.pdf\", recursive=True)\n","                if not pdf_files:\n","                    print(f\"No PDF found in {os.getcwd()}\")\n","                    os.chdir(\"..\")\n","                    os.rmdir(pmcid)\n","                    return False\n","\n","                pdf_file = pdf_files[0]\n","                print(f\"Found PDF: {pdf_file}\")\n","\n","                with open(pdf_file, 'rb') as local_pdf_file:\n","                    pdf_content = local_pdf_file.read()\n","\n","                with vfs.open(os.path.join(object_directory, f\"{pmcid}.pdf\"), 'wb') as vfs_pdf_file:\n","                    vfs_pdf_file.write(pdf_content)\n","\n","                print(f\"PDF {pdf_file} successfully written to TileDB VFS.\")\n","                os.chdir(\"..\")\n","                shutil.rmtree(pmcid)\n","                return True\n","        print(f\"No records found for PMCID {pmcid}.\")\n","        return False\n","    \n","    def fetch_abstract(pmid):\n","        \"\"\"Fetch the abstract from PubMed by PMID.\"\"\"\n","        url = \"https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi\"\n","        params = {'db': 'pubmed', 'id': pmid, 'retmode': 'xml'}\n","        file_name = os.path.join(object_directory, f\"{pmid}_abstract.txt\")\n","\n","        if vfs.is_file(file_name):\n","            print(f\"Abstract {pmid}_abstract.txt already exists.\")\n","            return True\n","        print(f\"Fetching abstract for PMID {pmid}...\")\n","        max_retries = 5\n","        retry_delay = 5\n","        for attempt in range(max_retries):\n","            try:\n","                response = requests.get(url, params=params)\n","                if response.status_code == 200:\n","                    root = ET.fromstring(response.content)\n","                    abstract_text = \"\\n\".join([abstract.text for abstract in root.findall(\".//AbstractText\")])\n","                    if not abstract_text.strip():\n","                        print(f\"No abstract available for PMID {pmid}.\")\n","                        return False\n","                    with vfs.open(file_name, 'wb') as file:\n","                        file.write(abstract_text.encode('UTF-8'))\n","                    print(f\"Abstract saved to {file_name}.\")\n","                    return True\n","\n","                elif response.status_code == 429:\n","                    print(f\"Rate limit hit for PMID {pmid}, retrying in {retry_delay} seconds...\")\n","                    time.sleep(retry_delay)\n","                    retry_delay *= 2\n","            except requests.exceptions.ConnectionError as e:\n","                print(f\"Network error on attempt {attempt + 1}: {e}\")\n","                time.sleep(retry_delay)\n","                retry_delay *= 2\n","        print(f\"Failed to retrieve abstract for PMID {pmid} after {max_retries} attempts.\")\n","        return False\n","\n","    # Add to file log function\n","    def add_to_new_file(file_name, prefix):\n","        \"\"\"Append the name of the processed file to the file history log in S3.\"\"\"\n","        new_file_path = f\"{object_directory}/history/{prefix}_{job_id}\"\n","        tmp_data = b\"\"\n","        try:\n","            if vfs.is_file(new_file_path):\n","                with vfs.open(new_file_path, 'rb') as f:\n","                    tmp_data = f.read()\n","\n","            file_name_as_bytes = (file_name + '\\n').encode('utf-8')\n","            tmp_data += file_name_as_bytes\n","\n","            with vfs.open(new_file_path, 'wb') as f:\n","                f.write(tmp_data)\n","        except Exception as e:\n","            print(f\"Error appending to file history: {e}\")\n","\n","    # Processing each entry in the DataFrame\n","    job_df.dropna(subset=['PMID Gene-disease'], inplace=True)\n","    total_entries = len(job_df)  \n","    for index, row in job_df.iterrows():\n","        try:\n","            pmid = str(int(row[\"PMID Gene-disease\"]))\n","            convert_pmid_to_pmcid_and_ingest(pmid)\n","        except ValueError as e:\n","            continue\n","\n","    total_time = time.time() - start_time\n","    print(f\"Processed {total_entries} entries in {total_time:.2f} seconds.\")\n"]},{"cell_type":"markdown","id":"95779637-e9ae-49a8-b384-08c9655feaaf","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}}},"source":["### Pipeline Factory Function\n","The `initialize_step` will build out our DAG starting with the ingestion steps. The below code: (write when it works)"]},{"cell_type":"code","execution_count":null,"id":"82981466-efde-43e1-829e-5eb689428bc4","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}},"trusted":true},"outputs":[],"source":["def pipeline_step(access_credentials: str, repo_url: str, out_directory: str, bucket_path: str,index_uri: str, in_file: str = \"rag-article-list.txt\", \n","                  num_jobs: int = 4, bucket_region: str = \"us-west-2\", hash_check: bool = True, model_name: str = \"pritamdeka/BioBERT-mnli-snli-scinli-scitail-mednli-stsb\",\n","                  index_type: str = \"ivf_flat_index\",dimensions: int = 768,ingestion_mode: str =\"BATCH\"):\n","    \"\"\"\n","    Pipeline step to pull the repo, check for changes using a hash, and divide the input file into jobs for batch processing.\n","\n","    :param access_credentials: Credentials to access TileDB Cloud for DAG execution.\n","    :param repo_url: Git repository URL to clone or pull the latest version from.\n","    :param out_directory: Output directory where the files will be stored.\n","    :param bucket_path: S3 bucket path for file storage.\n","    :param in_file: Input file name, defaulting to 'rag-article-list.txt'.\n","    :param num_jobs: Number of jobs to divide the input file into, default is 4.\n","    :param bucket_region: S3 bucket region, default is 'us-west-2'.\n","    :param hash_check: If True, the function checks if the file has changed using a hash before proceeding.\n","    :param index_uri: TileDB index URI where the embeddings are stored.\n","    :param model_name: Name of the embedding model to store in TileDB metadata.\n","    :param index_type: Type of index used for TileDB. Must be one of the supported types.\n","    :param input_files_dir: Directory from where the input files will be ingested.\n","    :param dimensions: The dimensionality of the embeddings.\n","    \"\"\"\n","\n","    # Supported index types\n","    supported_index_types = [\"FLAT\", \"IVF_FLAT\", \"IVF_PQ\", \"VAMANA\"]    \n","    # Validate index type\n","    if index_type not in supported_index_types:\n","        raise ValueError(f\"Unsupported index type '{index_type}'. Supported types are: {supported_index_types}\")\n","\n","    # Supported index ingestion types\n","    supported_ingestion_modes = [\"BATCH\",\"REALTIME\"]\n","    if ingestion_mode not in supported_ingestion_modes:\n","        raise ValueError(f\"Unsupported index_ingestion_mode type '{index_ingestion_mode}'. Supported types are: {ingestion_modes}\")\n","    \n","    # Step 1: Prepare paths\n","    full_bucket_path = f\"{bucket_path}/{out_directory}\"  # Path to the S3 bucket\n","\n","    # Step 2: Pull the latest changes from the repository\n","    home_directory = os.path.expanduser(\"~\")\n","    os.chdir(home_directory)\n","    \n","    repo_name = repo_url.split('/')[-1].replace('.git', '')  # Extract repo name from URL\n","\n","    if os.path.exists(repo_name):\n","        os.chdir(repo_name)\n","        subprocess.run([\"git\", \"pull\"])  # Pull the latest changes\n","    else:\n","        subprocess.run([\"git\", \"clone\", repo_url])  # Clone the repo if it doesn't exist\n","        os.chdir(repo_name)\n","\n","    # Step 3: Check and handle file hash\n","    previous_hash_path = \"previous_hash.txt\"\n","    current_hash_path = \"hash.txt\"\n","    current_hash = \"\"\n","\n","    if os.path.exists(current_hash_path):\n","        with open(current_hash_path, 'r') as f:\n","            current_hash = f.read().strip()\n","            print(f\"Current hash: {current_hash}\")\n","    else:\n","        if os.path.exists(in_file):\n","            current_hash = hash_file(in_file)  # Compute the hash of the input file\n","            with open(current_hash_path, 'w') as fh:\n","                fh.write(current_hash)\n","        else:\n","            print(f\"{in_file} does NOT exist. Please submit a valid file and rerun this function.\")\n","            return\n","\n","    # Step 4: Compare hash with previous hash\n","    if os.path.exists(previous_hash_path):\n","        with open(previous_hash_path, 'r') as f:\n","            previous_hash = f.read().strip()\n","            if current_hash == previous_hash:\n","                print(\"File did not change.\")\n","                if hash_check:\n","                    return  # Stop if hash_check is True and no changes\n","            else:\n","                print(\"File has changed, proceeding with processing.\")\n","                with open(previous_hash_path, 'w') as fh:\n","                    fh.write(current_hash)\n","    else:\n","        with open(previous_hash_path, 'w') as fh:\n","            fh.write(current_hash)\n","\n","    # Step 5: Commit the updated hash and file to the Git repository\n","    add_and_commit_files(\"Updating the hash and file to keep the repo up to date\")  # Commit changes\n","\n","    # Step 6: Load the DataFrame from the input file\n","    df = pd.read_csv(in_file, sep='\\t', engine='python')  # Load input file into a DataFrame\n","    total_rows = len(df)\n","    print(f\"Total rows in the DataFrame: {total_rows}\")\n","\n","    # Step 7: Divide the DataFrame into jobs based on num_jobs\n","    if num_jobs > total_rows:\n","        num_jobs = total_rows\n","        print(f\"Number of jobs reduced to {num_jobs} to match total rows.\")\n","\n","    job_size = math.ceil(total_rows / num_jobs)\n","    jobs = [df.iloc[i:i + job_size] for i in range(0, total_rows, job_size)]\n","    print(f\"Divided into {len(jobs)} jobs\")\n","\n","    # Step 8: Set up TileDB Cloud DAG for batch processing\n","    dag = tiledb.cloud.dag.DAG(name=\"document_batch\", mode=tiledb.cloud.dag.Mode.BATCH)\n","\n","    # Consolidation and deletion tasks\n","    consolidate_node = dag.submit(consolidate_chunks, num_jobs, bucket_region, full_bucket_path, \"file_history\", access_credentials_name=access_credentials)\n","    consolidate_node_2 = dag.submit(consolidate_chunks, num_jobs, bucket_region, full_bucket_path, \"missed_paper_history\", access_credentials_name=access_credentials)\n","    delete_unwanted_node = dag.submit(delete_unwanted_files, bucket_region, full_bucket_path, access_credentials_name=access_credentials)\n","    delete_unwanted_node.depends_on(consolidate_node)\n","\n","    # Step 9: Submit each job for ingestion\n","    for i, job_df in enumerate(jobs):\n","        if job_df.empty:\n","            print(\"Skipping empty job.\")\n","            continue\n","        cpu, mem = estimate_resources(job_df)\n","        print(f\"Processing {len(job_df)} rows, estimated CPU: {cpu}, estimated memory: {mem:.2f} GB.\")\n","\n","        ingest_node = dag.submit(\n","            pmid_ingestion,\n","            job_df,\n","            full_bucket_path,\n","            bucket_region,\n","            i,\n","            index_uri,\n","            model_name,\n","            access_credentials_name=access_credentials,\n","            image_name=\"vectorsearch\",\n","            resources={\"cpu\": f\"{cpu}\", \"memory\": f\"{mem}Gi\"}\n","        )\n","        consolidate_node.depends_on(ingest_node)\n","        consolidate_node_2.depends_on(ingest_node)\n","\n","    dag.compute()\n","    dag.wait()\n","    print(f\"Finished processing all jobs, consolidating history, and cleaning up undesired S3 files at {full_bucket_path}\")\n","\n","    # Step 10: Index creation and embedding ingestion\n","    embedding = SentenceTransformersEmbedding(model_name_or_path=model_name, dimensions=dimensions)\n","    reader = DirectoryTextReader(\n","        search_uri=full_bucket_path,\n","        include=\"*\",  \n","        suffixes=[\".pdf\", \".txt\"],\n","        exclude=[\"[.]*\", \"*/[.]*\"],  \n","        text_splitter=\"RecursiveCharacterTextSplitter\",\n","        text_splitter_kwargs={\"chunk_size\": 1000, \"chunk_overlap\": 100}\n","    )\n","\n","    # Clean the existing index\n","    print(\"Ensuring we have a clean index\")\n","    if tiledb.object_type(index_uri) == \"group\":\n","        print(\"Deleting existing index...\")\n","        group = tiledb.Group(index_uri, \"m\")\n","        group.delete(recursive=True)\n","        print(\"Deleted old index\")\n","\n","    print(\"Creating a new index\")\n","    index = object_index.create(\n","        uri=index_uri,\n","        index_type=index_type,\n","        object_reader=reader,\n","        embedding=embedding,\n","    )\n","    if ingestion_mode == \"BATCH\":\n","        embeddings_generation_driver_mode = tiledb.cloud.dag.Mode.BATCH\n","    elif ingestion_mode == \"REAL\":\n","        embeddings_generation_driver_mode = tiledb.cloud.dag.Mode.REAL\n","\n","    print(\"Kicking off the ingestion\")\n","    index.update_index(\n","        embeddings_generation_driver_mode=embeddings_generation_driver_mode,\n","        embeddings_generation_mode=embeddings_generation_driver_mode, #let's add onto this. \n","        files_per_partition=50, \n","        extra_worker_modules=[\"transformers==4.37.1\", \"PyMuPDF\", \"beautifulsoup4\"],\n","        embedding=embedding,\n","        worker_access_credentials_name=access_credentials,\n","        workers=num_jobs, # I believe we need to remove this if we choose REALTIME. I think this is maxed at 15. \n","        worker_resources={\"cpu\": f\"{cpu}\", \"memory\": f\"{mem}Gi\"}\n","    )\n","    \n","    return f\"Pipeline completed with file: {in_file}\"\n","\n","\n"]},{"cell_type":"code","execution_count":null,"id":"4a668b6d-c8ef-4eea-baf4-51fda6da6e34","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}},"trusted":true},"outputs":[],"source":["!pip install sentence-transformers"]},{"cell_type":"markdown","id":"dbecc262-8b80-48ac-86df-c698cb769764","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}}},"source":["### TileDB Ingestion Pipeline\n","The Pipeline can take a bit to run (depending on the size of the ingestion) so let's take a few moments to explore what is going on and you can read the code to learn more. "]},{"cell_type":"markdown","id":"ba4f0fa1-b708-42b0-bd6c-6604cc1eee5f","metadata":{},"source":["#### Key Features and Value Highlights of the Ingestion Pipeline Demo\n","\n","The ingestion pipeline has two major parts:\n","1. **PMC/PMCID Ingestion**: Scraping PubMed for documents using PMIDs to gather relevant medical articles and their abstracts.\n","2. **Multi-Modal TileDB Ingestion**: Ingesting the processed documents and data into TileDB, creating an index that allows for fast and scalable searches across multiple modalities (e.g., PDFs, text).\n","\n","This function highlights several key elements:\n","\n","#### Credential Management\n","One of the core benefits of the pipeline is how credentials and distributed tasks are abstracted across the entire process, simplifying the implementation and making it scalable.\n","\n","The `access_credentials` parameter allows for secure, centralized management of credentials, removing the need for hard-coded keys in scripts. This enables TileDB Cloud to handle requests automatically and securely.\n","\n","```python\n","def pipeline_step(access_credentials: str, repo_url: str, out_directory: str, bucket_path: str, index_uri: str, in_file: str = \"rag-article-list.txt\", \n","                  num_jobs: int = 4, bucket_region: str = \"us-west-2\", hash_check: bool = True, model_name: str = \"pritamdeka/BioBERT-mnli-snli-scinli-scitail-mednli-stsb\",\n","                  index_type: str = \"IFV_FLAT\", dimensions: int = 768, ingestion_mode: str =\"BATCH\"):\n","\n","```\n","\n","#### Scalable Vector Search and Indexing\n","\n","TileDB’s vector search capabilities allow for efficient search and indexing across datasets. The pipeline supports several index types, such as **IVF_FLAT**, **VAMANA**, and **IVF_PQ**, providing flexibility to optimize for either speed or memory efficiency.\n","\n","The ability to toggle between **batch (BATCH)** and **real-time (REALTIME)** ingestion modes offers flexibility to handle different scenarios. For batch processing, you can maximize throughput, while real-time ingestion enables rapid updates to the index.\n","\n","#### Unified Ingestion for Multi-Modal Data\n","\n","The pipeline supports **multi-modal data ingestion**, meaning it can handle PDFs, text files, and abstracts using the same loader. This eliminates the need for specialized logic for different data types. The unified loader handles different formats like PDFs and text seamlessly.\n","\n","In cases where documents are stored in compressed tarballs, the pipeline is capable of **recursive extraction** of PDFs, making it ideal for complex medical or academic data ingestion scenarios.\n","\n","#### Seamless Integration with TileDB Cloud DAG\n","\n","TileDB’s **Cloud DAG** allows you to manage complex workflows with distributed jobs like consolidating chunks and data cleanup. The DAG framework ensures tasks are executed in the correct order while supporting parallel processing and efficient scaling.\n","\n","This feature simplifies managing dependencies between tasks, such as ensuring that data cleanup only occurs after ingestion has been completed.\n","\n","#### Efficient Use of TileDB's Virtual File System (VFS)\n","\n","The pipeline leverages **TileDB's Virtual File System (VFS)** to interact directly with object stores like S3, without the need for local storage. This results in faster data transfers and reduces the storage footprint during the ingestion process.\n","\n","Additionally, the VFS allows tracking of file history and logs in S3, enabling easier debugging, auditing, and progress monitoring without needing to rely on local files.\n","\n","#### Support for Embedding Models and Metadata Storage\n","\n","The pipeline integrates seamlessly with embedding models like **BioBERT** or **HuggingFace**, allowing for the storage of embeddings alongside original documents. This is crucial for **retrieval-augmented generation (RAG)** systems or any pipeline that requires similarity-based search.\n","\n","The ability to store model metadata alongside the embeddings ensures **traceability**, allowing for reproducibility when running searches or processing documents in future versions of the pipeline.\n","\n","##### Why Experimenting with an Embedding Model is Important\n","\n","Experimenting with medical-specific embedding models like BioBERT is vital because it ensures the embeddings capture domain-specific context, terminology, and nuances. This precision enhances the quality of responses from LLMs when generating or retrieving information from scientific papers. By tying embeddings to the LLM's response, the system delivers more accurate, contextually relevant answers, crucial for applications in medical research or clinical insights where even slight misinterpretations can lead to incorrect conclusions.  Another way to think about this is in the context of \"human in the loop\" scenarios where humans label complex datasets. Just as medical embedding models require domain-specific training to understand intricate concepts, human labelers also need specialized knowledge to correctly annotate data. For example, labeling medical datasets is far more complex than identifying basic objects like \"this is a bird.\" If a non-expert incorrectly labels a medical condition, it could misinform the model, leading to inaccurate predictions or insights. Embedding models, like humans in these tasks, need to be properly trained with the right context to ensure meaningful and accurate representations.\n","\n","#### Real-Time Feedback with Logging and File History\n","\n","A key strength of the pipeline is its ability to provide **real-time feedback** on the progress of document ingestion. By logging successfully ingested files and noting failed attempts (through `file_history` and `missed_paper_history`), the pipeline allows you to quickly diagnose and address issues.\n","\n","This feedback loop ensures that failed files can be reprocessed without having to re-run the entire ingestion pipeline. Logs are stored directly in the object store for easy retrieval and auditing.\n","\n","#### Flexible Configuration and Extensibility\n","\n","The pipeline is highly configurable, with parameters such as:\n","\n","- **Number of jobs**: To control workload distribution.\n","- **Bucket path**: To specify object store locations.\n","- **Index type**: To optimize for different search strategies (e.g., IVF_FLAT for faster searches or VAMANA for memory efficiency).\n","- **Embedding dimensions**: To configure the model’s vector size (e.g., 768 for BioBERT).\n","\n","This flexibility allows the pipeline to scale based on dataset size and processing needs. Furthermore, the ability to switch between **batch** and **real-time** modes allows for dynamic processing depending on the workload.\n","\n","#### Summary\n","\n","This ingestion pipeline leverages **distributed abstraction**, **secure credential handling**, and **scalable vector search capabilities**, while remaining flexible for multi-modal data ingestion. The built-in features such as **change detection**, **retry logic**, and **embedding model integration** make this pipeline robust and scalable, handling even the most complex medical document ingestion and processing tasks.\n"]},{"cell_type":"markdown","id":"02596261-ce73-4715-b50c-1b91dc08882b","metadata":{},"source":["## Idexing and Ingestion Exploration and Guidelines \n","### Driver and Embedding Generation\n","For a deeper exploration, during the indexing/ingestion of the PDFs and text we see two DAGs being created. One of them is the ** *ingest_embeddings_with_driver task.***\n","\n","The ingest_embeddings_with_driver function, specifically through the driver, is responsible for orchestrating and scheduling tasks to the workers (pods) by setting up and managing the task execution flow. \n","\n","#### Workers and Ingestion Tasks\n","\n","* **Task Distribution:** The total set of files (or partitions) is divided into smaller groups based on object_partitions_per_worker. These groups are then assigned to workers, with each worker processing a specific number of partitions in parallel.\n","\n","* **Workers Allocation:** The number of workers (pods) is determined by the workers variable. Tasks are distributed across the workers, and if there are fewer tasks than workers, some workers will remain idle.\n","\n","*  **Dynamic Worker Usage:**  If there are fewer partitions than workers, only the necessary workers will be used, and any excess workers will not be spun up. If there are more partitions than workers, the system will continue assigning tasks to available workers until all partitions are processed."]},{"cell_type":"markdown","id":"0aaf8333-abfa-4de3-9492-afdc902911c8","metadata":{},"source":["### Task Scheduling Guidelines \n","When optimizing ingestion tasks in a distributed environment (e.g., Kubernetes), consider the following guidelines to maximize parallelism and resource efficiency:\n","\n","* **Task Size:** Keep ingestion tasks small enough to fit comfortably within a pod’s memory and CPU limits. Ideally, each task should process one file or partition at a time to reduce memory footprint and avoid resource contention.\n","\n","* **Memory Footprint:** Assign memory just sufficient for loading and processing a single file. This prevents over-allocation of resources while ensuring that each pod can handle its workload efficiently. Adjust memory requests/limits based on file size and processing complexity.\n","\n","* **Parallelism:** Use smaller pods with fewer tasks to maximize parallelism. More, smaller pods allow Kubernetes to distribute tasks across multiple CPU cores or nodes, improving throughput.\n","\n","* **Resource Allocation:** Ensure that each pod’s resource requests (CPU, memory) are tuned to its task size. Smaller resource requests lead to better scheduling flexibility and faster execution.\n","\n","By adhering to these guidelines, you can achieve better resource utilization, faster task execution, and improved scalability across your infrastructure."]},{"cell_type":"markdown","id":"5e0f7532-b770-4852-8d5b-3828e0ea525f","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}}},"source":["## Optional Phase: Download the consolidated file from S3 and write it to your local Git repository\n","let's also do this for our missed page. "]},{"cell_type":"code","execution_count":null,"id":"d6525fcd-dd25-4414-8f7c-5e50e8464bc9","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}},"trusted":true},"outputs":[],"source":["# Optional: Download the consolidated file from S3 and write it to your local Git repository\n","# Set up TileDB context with S3 credentials (replace with your actual credentials)\n","#ctx = tiledb.Ctx({\n"," #   \"vfs.s3.region\": \"\",  # S3 region\n","  #  \"vfs.s3.aws_access_key_id\": \"\",  # AWS Access Key (add your key here)\n","# \"vfs.s3.aws_secret_access_key\": \"\"  # AWS Secret Access Key (add your key here)\n","#})\n","\n","# S3 path to the consolidated file\n","#consolidated_file_path = \"\"\n","\n","# Initialize TileDB VFS to interact with S3\n","#vfs = tiledb.VFS(ctx=ctx)\n","\n","# Step 1: Read the contents of the consolidated file from S3\n","#file_contents = \"\"\n","#with vfs.open(consolidated_file_path, 'rb') as f:  # Open the file in binary read mode\n","    #file_contents = f.read()  # Read the entire file contents\n","\n","# Step 2: Write the contents to the local file system in the current directory\n","# You may want to navigate to your Git repository directory before writing\n","# Uncomment and modify the following line if needed:\n","# os.chdir(\"pmc-llm\")  # Change to your Git repository directory\n","\n","# Create the local path for the consolidated file\n","#path = os.path.join(os.getcwd(), \"consolidated_file_history\")  # Store in the current working directory\n","\n","# Write the contents from S3 to the local file in binary write mode\n","#with open(path, 'wb') as f:\n","    #f.write(file_contents)  # Write the downloaded file contents to the local file\n","\n","# Step 3: Assuming `add_and_commit_files` is a valid function, commit the changes to Git\n","# This function should add the file, commit it with a message, and push the changes\n","#add_and_commit_files(\"updating the consolidated file history for tracking purposes\")\n"]},{"cell_type":"markdown","id":"2577f8fb-85af-4135-b22c-1619ff2c4d79","metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","views":{"grid_default":{"col":null,"height":2,"hidden":true,"locked":true,"row":null,"width":2}}}}},"source":["## Optional Phase: Exploring our Ingested Files (WIP)\n","Every array registered with TileDB Cloud must be accessed using a URI of the form  `tiledb://<namespace>/<array-name>`, where `<namespace>` is the user or organization who owns the array and <array-name> is the array name set by the owner upon array registration. This URI is [displayed on the console when viewing the array details](https://docs.tiledb.com/cloud/how-to/arrays/view-array-details).\n","Set the TileDB configuration parameters rest.username and rest.passwordwith your TileDB Cloud username and password, or alternatively [rest.token with the API token you created.](https://docs.tiledb.com/cloud/how-to/account/create-api-tokens)  **Accessing arrays by setting an API token is typically faster than using your username and password.** More details [here](https://docs.tiledb.com/cloud/api-reference/array-access) First, let's set up som variables and explore some outputs\n"]},{"cell_type":"markdown","id":"cef8ff5b-c5d5-4ebd-9b20-4bec45d5f4c5","metadata":{},"source":["We can then view the schema of the array itself to better understand our vectors. "]},{"cell_type":"markdown","id":"cdfd731f-590b-4e5e-bb3a-f0e24807387a","metadata":{},"source":["TileDB stores the data in an efficient, compressed format, and when you query it, the data is returned as a 1D array to minimize storage complexity and retrieval overhead. "]},{"cell_type":"markdown","id":"6e4ec894-2345-4fd4-b51e-945e81ee27b7","metadata":{},"source":["We can also then figure out the total expected size of our array of vectors based on the non empty domains representing partitions. "]},{"cell_type":"markdown","id":"26a30a89-b422-4e83-bf55-a0b7c60f527d","metadata":{},"source":["| Algorithm         | Description                               | Accuracy | Speed  | Memory Efficiency | When to Use                                            |\n","|-------------------|-------------------------------------------|----------|--------|-------------------|--------------------------------------------------------|\n","| **flat_index**     | FlatIndex (Brute-force search)            | Highest  | Slow   | Low               | Small datasets where accuracy is the top priority       |\n","| **ivf_flat_index** | IVFFlat Index (Inverted File with Flat)   | Medium   | Fast   | Medium            | Large-scale datasets, good balance between speed and accuracy |\n","| **vamana_index**   | Vamana Index (Graph-based search)         | High     | Fast   | Low               | When high accuracy is required with faster search times compared to flat_index |\n","| **ivf_pq_index**   | IVFPQ Index (Inverted File with PQ)       | Low      | Very Fast | High            | Very large datasets, where memory efficiency is a priority over accuracy |\n"]},{"cell_type":"markdown","id":"866871cd-f029-4376-bdb6-68fc388c349f","metadata":{},"source":["Each partition ID in your TileDB array points to a batch of 258 arrays, and each of those arrays (or vectors) has 768 float32 values. The external ID serves as a unique identifier for each vector within that batch.\n","\n","So, to summarize:\n","\n","* Partition ID: This is the key that points to a \"cell\" in the TileDB array.\n","* Batch of 258 Vectors: Each cell contains 258 vectors, where each vector represents one embedding or feature array of 768 float32 values.\n","* External ID: For every vector in the batch, there's an external ID that uniquely maps to that specific vector. The external ID helps link or identify which specific vector (or array) the data corresponds to within the batch.\n","\n","You have 7 partition IDs (from 0 to 6).\n","Each partition contains 258 vectors.\n","So, the total number of vectors in your TileDB array is 1,806 vectors. \n","\n","Example Scenario:\n","    * Partition contains: 258 vectors (batch).\n","    * Query: You want to find the top 10 nearest neighbors (k=10).\n","    * Result: You load one partition (258 vectors), compute the similarity between your query vector and all 258 vectors, and extract the top 10 most similar vectors. Since k=10 is less than 258, you don't need to fetch another partition, making this operation highly efficient.\n","    \n","When You Might Need More Tiles:\n","If k > 258, you may need to fetch additional partitions to get enough vectors for the query, but with good locality, you can still minimize the number of partitions you need to pull.\n","\n","\n","Summary:\n","* Locality advantage: Similar vectors are likely stored close together in the same partition, so once you find one similar vector, you can search the entire batch efficiently.\n","* Efficient for small k: If k < 258, you can find all your nearest neighbors within a single partition, reducing the need for additional I/O.\n","* Improved performance: Fewer disk reads, fewer memory loads, and better cache utilization make the search faster and more efficient.\n"]},{"cell_type":"markdown","id":"2efeb70c-19f8-45ee-ac04-17066f13074e","metadata":{},"source":["## Hugging Face Pipelines\n","a Hugging Face pipeline provides a high-level API that simplifies the entire process of loading models and performing tasks such as text generation, sentiment analysis, translation, and more. It abstracts away the complexities of:\n","\n","    Loading the model: Automatically retrieves and initializes the model from Hugging Face’s model hub.\n","    Tokenization: Tokenizes the input text before passing it to the model.\n","    Running inference: Submits the tokenized input to the model for processing.\n","    Post-processing: Converts model outputs back into human-readable format.\n","\n","This makes it easier to work with NLP models without needing to manually handle every step.\n","\n","## LLM Response\n","Now that we've explored our embeddings, its time to get an LLM response. This section of the code can and should be executed on a GPU notebook. Luckily, with TileDB you saved code can always access the data via TileDB URL and our authentication methods, so you can choose what type of notebooks to execute or iterate on various task. All the above task are using task graphs and standard Python, so do not benefit from GPUs. The Tasks below will use a GPU to improve response times. "]},{"cell_type":"code","execution_count":null,"id":"92019157-ca14-4cbb-b8cd-2dd7a78102f2","metadata":{"trusted":true},"outputs":[],"source":["import transformers\n","from transformers import AutoTokenizer, pipeline\n","from langchain.llms import HuggingFacePipeline\n","from langchain.prompts import PromptTemplate\n","from langchain.chains import LLMChain\n","from langchain.vectorstores.tiledb import TileDB\n","import tiledb.cloud\n","from tiledb.vector_search.object_api import object_index\n","\n","# Function to initialize the LLM chain\n","def initialize_llmchain(llm_model, temperature, max_tokens):\n","    tokenizer = AutoTokenizer.from_pretrained(llm_model)\n","    from transformers import BitsAndBytesConfig\n","    quantization_config = BitsAndBytesConfig(load_in_4bit=True,\n","                                         llm_int8_threshold=200.0)\n","    hf_pipeline = transformers.pipeline(\n","        \"text-generation\",\n","        model=llm_model,\n","        model_kwargs={\n","            'quantization_config': quantization_config,\n","        },\n","        tokenizer=tokenizer,\n","        trust_remote_code=True,\n","        device_map=\"auto\",\n","        torch_dtype=\"bfloat16\",\n","        max_new_tokens=max_tokens,\n","        eos_token_id=tokenizer.eos_token_id,\n","        pad_token_id=tokenizer.eos_token_id,\n","    )\n","    llm = HuggingFacePipeline(pipeline=hf_pipeline, model_kwargs={'temperature': temperature})\n","\n","    # Define the prompt template\n","    template = \"\"\"You are a helpful assistant. \n","Context:\n","{context}\n","\n","Conversation:\n","{chat_history}\n","User: {question}  \n","Use the context information only if related to the question.\n","Generate only one assistant response.\n","Output:\"\"\"\n","    prompt = PromptTemplate.from_template(template)\n","    llm_chain = LLMChain(prompt=prompt, llm=llm)\n","\n","    return llm_chain\n","\n","# Initialize the LLM and Vector Search\n","def initialize_LLM(index_uri, llm_name, llm_temperature, max_tokens):\n","    index = object_index.ObjectIndex(index_uri, memory_budget=-1)\n","    qa_chain = initialize_llmchain(llm_name, llm_temperature, max_tokens)\n","    return index, qa_chain\n","\n","# Format chat history for display\n","def format_chat_history(message, chat_history):\n","    recent_history = chat_history[-10:] if len(chat_history) > 3 else chat_history\n","    formatted_chat_history = \"\"\n","    for user_message, bot_message in recent_history:\n","        formatted_chat_history += f\"User: {user_message}\\n\"\n","        formatted_chat_history += f\"Assistant: {bot_message}\\n\"\n","    return formatted_chat_history\n","\n","# Perform conversation with LLM and VectorDB\n","def conversation(vector_db, topk, qa_chain, message, history):\n","    formatted_chat_history = format_chat_history(message, history)\n","    question = message\n","    _, _, results = vector_db.query(\n","        {\"text\": [question]}, \n","        k=topk,\n","        nprobe=vector_db.index.partitions,\n","        return_objects=False,\n","        return_metadata=True,\n","    )\n","\n","    # Build context from the retrieved documents\n","    context = \"\"\n","    for i, text in enumerate(results[\"text\"][0]):\n","        context += f\"Context for file with path: {results['file_path'][0][i]} at page: {results['page'][0][i]}\\n{text.strip()}\\n\"\n","    \n","    # Print the retrieved documents for debugging\n","  #  print(\"\\nRetrieved Documents:\")\n","   # for doc in results[\"text\"][0]:\n","    #    print(doc)\n","    #print path from the retrieved document. \n","    print(\"\\nFile Paths:\")\n","    for path in results[\"file_path\"][0]:\n","        print(path)\n","    # Generate response from LLM\n","    response_answer = qa_chain.run(question=question, context=context, chat_history=formatted_chat_history)\n","    \n","    # Append to chat history\n","    new_history = history + [(message, response_answer)]\n","    \n","    return vector_db, qa_chain, new_history\n"]},{"cell_type":"code","execution_count":null,"id":"2f33c769-97e8-478d-bcc1-0038d275d568","metadata":{"trusted":true},"outputs":[],"source":["!pip install gradio\n","!pip install transformers --upgrade\n","!pip install torch --upgrade"]},{"cell_type":"code","execution_count":null,"id":"b36a7ef3-62dc-48f7-9b16-b8181652ef9b","metadata":{"trusted":true},"outputs":[],"source":["import gradio as gr  # Ensure Gradio is imported\n","def llm_call(query: str, index_directory: str, max_tokens: int, top_k: int, llm_temperature: float):\n","    import transformers\n","    from transformers import AutoTokenizer, pipeline\n","    from langchain.llms import HuggingFacePipeline\n","    from langchain.prompts import PromptTemplate\n","    from langchain.chains import LLMChain\n","    from langchain.vectorstores.tiledb import TileDB\n","    import tiledb.cloud\n","    from tiledb.vector_search.object_api import object_index\n","    \n","    # Get the vector DB URI\n","    user_profile = tiledb.cloud.user_profile()\n","    index_uri = f\"tiledb://{user_profile.username}/{user_profile.default_s3_path.rstrip('/')}{index_directory}\"\n","    llm_name = \"BioMistral/BioMistral-7B\"\n","    \n","    # Initialize LLM Chain\n","    def initialize_llmchain(llm_model, temperature, max_tokens, top_k):\n","        tokenizer = AutoTokenizer.from_pretrained(llm_model)\n","        from transformers import BitsAndBytesConfig\n","        quantization_config = BitsAndBytesConfig(load_in_4bit=True, llm_int8_threshold=200.0)\n","        \n","        hf_pipeline = transformers.pipeline(\n","            \"text-generation\",\n","            model=llm_model,\n","            model_kwargs={'quantization_config': quantization_config},\n","            tokenizer=tokenizer,\n","            trust_remote_code=True,\n","            device_map=\"auto\",\n","            torch_dtype=\"bfloat16\",\n","            max_new_tokens=max_tokens,\n","            eos_token_id=tokenizer.eos_token_id,\n","            pad_token_id=tokenizer.eos_token_id,\n","        )\n","        llm = HuggingFacePipeline(pipeline=hf_pipeline, model_kwargs={'temperature': temperature})\n","\n","        # Define the prompt template\n","        template = \"\"\"You are a helpful assistant. \n","        Context:\n","        {context}\n","        \n","        Conversation:\n","        {chat_history}\n","        User: {question}  \n","        \n","        Use the context information only if related to the question.\n","        Generate only one assistant response.\n","        Output:\"\"\"\n","        \n","        prompt = PromptTemplate.from_template(template)\n","        llm_chain = LLMChain(prompt=prompt, llm=llm)\n","        return llm_chain\n","\n","    # Initialize LLM and Vector Search\n","    def initialize_LLM(index_uri, llm_name, llm_temperature, max_tokens, top_k):\n","        index = object_index.ObjectIndex(index_uri, memory_budget=-1)\n","        qa_chain = initialize_llmchain(llm_name, llm_temperature, max_tokens, top_k)\n","        return index, qa_chain\n","\n","    # Format chat history\n","    def format_chat_history(message, chat_history):\n","        recent_history = chat_history[-10:] if len(chat_history) > 3 else chat_history\n","        formatted_chat_history = \"\"\n","        for user_message, bot_message in recent_history:\n","            formatted_chat_history += f\"User: {user_message}\\n\"\n","            formatted_chat_history += f\"Assistant: {bot_message}\\n\"\n","        return formatted_chat_history\n","\n","    # Perform conversation\n","    def conversation(vector_db, topk, qa_chain, message, history):\n","        formatted_chat_history = format_chat_history(message, history)\n","        question = message\n","        _, _, results = vector_db.query(\n","            {\"text\": [question]}, \n","            k=topk,\n","            nprobe=vector_db.index.partitions,\n","            return_objects=False,\n","            return_metadata=True,\n","        )\n","\n","        # Build context from retrieved documents\n","        context = \"\"\n","        for i, text in enumerate(results[\"text\"][0]):\n","            context += f\"Context for file with path: {results['file_path'][0][i]} at page: {results['page'][0][i]}\\n{text.strip()}\\n\"\n","\n","        # Collect unique file paths into a set\n","        file_paths_set = set(results[\"file_path\"][0])\n","\n","        # Combine the file paths into a single string\n","        combined_file_paths = \"\\n\".join(file_paths_set)\n","\n","        # Generate response from LLM\n","        response_answer = qa_chain.run(question=question, context=context, chat_history=formatted_chat_history)\n","\n","        # Append to chat history\n","        new_history = history + [(message, response_answer)]\n","        return vector_db, qa_chain, new_history, combined_file_paths\n","\n","    # Initialize LLM and Vector Search\n","    vector_db, qa_chain = initialize_LLM(index_uri, llm_name, llm_temperature, max_tokens, top_k)\n","\n","    # Example conversation\n","    history = []\n","    vector_db, qa_chain, history, combined_file_paths = conversation(vector_db, top_k, qa_chain, query, history)\n","\n","    # Format final response for display\n","    final_response = f\"Retrieved Files:\\n{combined_file_paths}\\n\\n\"\n","    for user_msg, assistant_msg in history:\n","        final_response += f\"User: {user_msg}\\n\\nAssistant: {assistant_msg}\\n\\n\"\n","    \n","    return final_response\n"]},{"cell_type":"code","execution_count":null,"id":"360d3ffb-9d14-4800-8559-464311ddc8d3","metadata":{"trusted":true},"outputs":[],"source":["def fetch_and_display_pdf(aws_access_key_id, aws_secret_access_key, s3_region, s3_url):\n","    # Set up the TileDB VFS context with the provided credentials\n","    ctx = tiledb.Ctx({\n","        \"vfs.s3.region\": s3_region,\n","        \"vfs.s3.aws_access_key_id\": aws_access_key_id,\n","        \"vfs.s3.aws_secret_access_key\": aws_secret_access_key\n","    })\n","    \n","    # Extract the bucket and path from the s3_url\n","    bucket_name = s3_url.split('/')[2]\n","    object_path = '/'.join(s3_url.split('/')[3:])\n","    \n","    # Set up TileDB VFS to read the PDF\n","    vfs = tiledb.VFS(ctx=ctx)\n","    \n","    # Temporary file to store the PDF locally\n","    local_pdf_path = \"/tmp/fetched_pdf.pdf\"\n","    \n","    # Fetch the PDF from S3\n","    with vfs.open(f\"s3://{bucket_name}/{object_path}\", \"rb\") as remote_pdf:\n","        with open(local_pdf_path, \"wb\") as local_pdf:\n","            local_pdf.write(remote_pdf.read())\n","    \n","    # Now return the PDF file path for visualization in Gradio\n","    return local_pdf_path"]},{"cell_type":"code","execution_count":null,"id":"736706c3-3e3a-46ed-ba61-2acafaf56c4f","metadata":{"trusted":true},"outputs":[],"source":["jupyterhub_user = os.getenv(\"JUPYTERHUB_USER\")\n","jupyterhub_api_url = \"https://us-west-2.aws.jupyterhub.cloud.tiledb.com\"\n","gradio_port = 7830\n","gradio_path = f\"/user/{jupyterhub_user}/proxy/{gradio_port}/\"\n","gradio_frontend_url=f\"{jupyterhub_api_url}{gradio_path}\"\n","# Define the Gradio interface\n","def demo():\n","\n","    print(f\"JupyterHub User = {jupyterhub_user} jupyterhub_api_url = {jupyterhub_api_url} gradio port = {gradio_port} gradio path = {gradio_path} gradio_frontend_url = {gradio_frontend_url}\")\n","    with gr.Blocks() as interface:\n","        with gr.Tab(\"Ingestion Pipeline\"):\n","            gr.Markdown(\"### Enter the details for the ingestion pipeline.\")\n","\n","            in_file = gr.Textbox(label=\"Enter File Name (e.g., 'data.txt')\", value=\"rag-article-list.txt\")\n","            access_credentials = gr.Textbox(label=\"Access Credentials\", value=\"\")\n","            repo_url = gr.Textbox(label=\"Repo URL\", value=\"https://github.com/TileDB-Inc/pmc-llm.git\")\n","            out_directory = gr.Textbox(label=\"Out Directory\", value=\"pmc/rag/ingestion\")\n","            bucket_path = gr.Textbox(label=\"Bucket Path\", value=\"\")\n","            model_name = gr.Textbox(label=\"Model Name\", value=\"\")\n","            index_uri = gr.Textbox(label=\"Index URI\", value=\"\")\n","            num_jobs = gr.Number(label=\"Number of Jobs\", value=4)\n","            bucket_region = gr.Textbox(label=\"Bucket Region\", value=\"\")\n","            hash_check = gr.Checkbox(label=\"Enable Hash Check\", value=False)\n","            index_type = gr.Textbox(label=\"Index Type\", value=\"IVF_FLAT\")\n","            dimensions = gr.Number(label=\"Embedding Dimensions\", value=768)\n","            ingestion_output = gr.Textbox(label=\"Ingestion Output\")\n","\n","            ingest_button = gr.Button(\"Start Ingestion\")\n","            ingest_button.click(\n","                fn=pipeline_step,\n","                inputs=[access_credentials, repo_url, out_directory, bucket_path, index_uri, in_file, num_jobs, bucket_region, hash_check, model_name, index_type, dimensions],\n","                outputs=ingestion_output\n","            )\n","\n","        with gr.Tab(\"LLM Query\"):\n","            gr.Markdown(\"### Query the LLM with a custom question.\")\n","            \n","            query_input = gr.Textbox(label=\"Query\", value=\"What are the effects of halothane on the fetus?\")\n","            index_directory = gr.Textbox(label=\"Index Directory\", value=\"/pmc/rag/index\")\n","            max_tokens = gr.Slider(label=\"Max Tokens\", minimum=100, maximum=1000, step=100, value=500)\n","            top_k = gr.Slider(label=\"Top K\", minimum=1, maximum=10, step=1, value=3)\n","            llm_temperature = gr.Slider(label=\"LLM Temperature\", minimum=0.1, maximum=1.0, step=0.1, value=0.7)\n","            llm_output = gr.Textbox(label=\"LLM Output\")\n","\n","            query_button = gr.Button(\"Query LLM\")\n","            query_button.click(\n","                fn=llm_call,\n","                inputs=[query_input, index_directory, max_tokens, top_k, llm_temperature],\n","                outputs=llm_output\n","            )\n","        with gr.Tab(\"PDF Viewer\"):\n","            gr.Markdown(\"### Enter AWS credentials and S3 URL to fetch the PDF.\")\n","\n","            aws_access_key_id = gr.Textbox(label=\"AWS Access Key ID\", value=\"\")\n","            aws_secret_access_key = gr.Textbox(label=\"AWS Secret Access Key\", value=\"\", type=\"password\")\n","            s3_region = gr.Textbox(label=\"S3 Region\", value=\"us-west-2\")\n","            s3_url = gr.Textbox(label=\"S3 URL\", value=\"s3://your-bucket/path-to-pdf.pdf\")\n","\n","            pdf_output = gr.File(label=\"Fetched PDF\")\n","\n","            fetch_button = gr.Button(\"Fetch PDF\")\n","            fetch_button.click(\n","                fn=fetch_and_display_pdf,\n","                inputs=[aws_access_key_id, aws_secret_access_key, s3_region, s3_url],\n","                outputs=pdf_output\n","            )            \n","        interface.launch(\n","            debug=False,\n","            quiet=True,\n","            show_error=True,\n","            height=1100,\n","            root_path=\"/user/9a6ff6a4-7b46-490c-93eb-bd067509ac74/proxy/7830/\",  # Adjust for your proxy path\n","            server_name=\"0.0.0.0\",  # Bind to all interfaces, not just localhost\n","            server_port=7830  # Ensure you're using the correct port\n","        )\n","\n"]},{"cell_type":"code","execution_count":null,"id":"db7e8393-650b-4e3e-b7dd-0b1d2f285762","metadata":{"trusted":true},"outputs":[],"source":["# Run the Gradio interface\n","if __name__ == \"__main__\":\n","    demo()"]},{"cell_type":"code","execution_count":null,"id":"6760b9cb-41b4-4e03-b23e-9b650ef479d4","metadata":{"trusted":true},"outputs":[],"source":["from IPython.display import display, IFrame\n","print(f\"gradio_frontend_url= {gradio_frontend_url}\")\n","IFrame(gradio_frontend_url, width=\"100%\", height=\"1100\")"]}],"metadata":{"extensions":{"jupyter_dashboards":{"activeView":"grid_default","version":1,"views":{"grid_default":{"cellMargin":2,"defaultCellHeight":40,"maxColumns":12,"name":"grid","type":"grid"}}}},"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.9.20"}},"nbformat":4,"nbformat_minor":5}