part_DIY_RAG_pdf.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f1d713b3-579f-4ee2-93fd-541caf94e75c",
   "metadata": {
    "tags": []
   },
   "source": [
    "##### This Notebook use Doc AI and converts all PDF document to text so that down the line process can work\n",
    "\n",
    "#### Author: Saurabh Mangal (saurabhmangal@google.com)\n",
    "#### Editor / Reviewer: Wan Qi, Jing Le, Renzo Garcia\n",
    "##### Date: 11th Mar\n",
    "##### Description: This notebook contains part 1 of lab\n",
    "\n",
    " Copyright (c) [2024] [saurabhmangal@] -- \n",
    " This notebook is licensed under the Commercial License."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82f94fab-751f-4861-b06f-502a9c458e56",
   "metadata": {},
   "source": [
    "### Install Dependencies & Restart Kernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23d4a190-bfb9-4638-85c2-05f87b4c087a",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "%pip install --quiet PyPDF2\n",
    "%pip install --quiet pdfreader\n",
    "!pip install --quiet google-cloud-discoveryengine\n",
    "!pip install --upgrade --quiet google-cloud-storage\n",
    "!pip3 install --upgrade --quiet google-cloud-documentai\n",
    "!pip3 install --upgrade --quiet google-cloud-storage\n",
    "!pip3 install --upgrade --quiet google-cloud-documentai-toolbox\n",
    "%pip install --upgrade --quiet  google-cloud-documentai\n",
    "%pip install --upgrade --quiet google-cloud-aiplatform -q\n",
    "%pip install tqdm -q\n",
    "%pip install langchain -q\n",
    "!pip install --upgrade --quiet  langchain-google-genai\n",
    "%pip install --upgrade --quiet  google-cloud-documentai\n",
    "%pip install --upgrade --quiet  google-cloud-documentai-toolbox\n",
    "%pip install --upgrade --quiet  langchain-core langchain-google-vertexai\n",
    "\n",
    "#Restart kernel after installs so that your environment can access the new packages\n",
    "import IPython\n",
    "import time\n",
    "\n",
    "app = IPython.Application.instance()\n",
    "app.kernel.do_shutdown(True)\n",
    "\n",
    "time.sleep(10)\n",
    "\n",
    "print(\"Installation done\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86e4bb19-4949-47c5-b15a-b92558859c7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Run this cell and all below "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87d0b933-ba60-4434-8030-2e6a7279546f",
   "metadata": {},
   "source": [
    "### Declare Variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "924b9b34-af9d-4af3-b42e-c93e92482ad5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#no spaces or special characters allowed), ensure that it is unique\n",
    "import socket\n",
    "import re\n",
    "\n",
    "UNIQUE_PREFIX = socket.gethostname()\n",
    "\n",
    "UNIQUE_PREFIX = re.sub('[^A-Za-z0-9]+', '', UNIQUE_PREFIX)\n",
    "\n",
    "PREFIX_NUM_ONLY=int(str(re.search(r'\\d+', UNIQUE_PREFIX).group()))\n",
    "\n",
    "REGION_ALLOCATE=PREFIX_NUM_ONLY%3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "894d7f13-a452-4945-a9d0-73fef0c30af0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "PROJECT_ID = !(gcloud config get-value core/project)\n",
    "PROJECT_ID = PROJECT_ID[0]\n",
    "\n",
    "SVC_ACC = !(gcloud config get-value core/account)\n",
    "SVC_ACC = SVC_ACC[0]\n",
    "\n",
    "PROJECT_NUMBER=str(re.search(r'\\d+', SVC_ACC).group())\n",
    "\n",
    "LOCATION=\"asia-southeast1\"\n",
    "\n",
    "FOLDER_NAME=\".\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c39e8f13-13f0-409f-9a6a-d097a9a82ae6",
   "metadata": {},
   "source": [
    "### Create GCS Bucket & Import Documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f2b831f-a2ff-4c0d-9db2-47ba4a1938c0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from google.cloud import storage\n",
    "\n",
    "client = storage.Client()\n",
    "\n",
    "GCS_BUCKET_LOCATION = \"asia-southeast1\"\n",
    "\n",
    "GCS_BUCKET_NAME = f\"{PROJECT_ID}-{UNIQUE_PREFIX}\"\n",
    "GCS_BUCKET_URI = f\"gs://{GCS_BUCKET_NAME}\"\n",
    "\n",
    "bucket = storage.Bucket(client, GCS_BUCKET_NAME)\n",
    "\n",
    "if bucket.exists()==False:\n",
    "    # Create a Cloud Storage Bucket\n",
    "    !gcloud storage buckets create $GCS_BUCKET_URI --location=$GCS_BUCKET_LOCATION\n",
    "\n",
    "    # Upload the PDFs located in the books/ directory into the GCS bucket that you created\n",
    "    !gsutil cp -r $FOLDER_NAME/books/* $GCS_BUCKET_URI/books\n",
    "\n",
    "    # Verify that all Books 1 to 7 are uploaded to the GCS bucket (8 files in total, 2 for Part 1)\n",
    "    !gsutil ls $GCS_BUCKET_URI\n",
    "else:\n",
    "    # Upload the PDFs located in the books/ directory into the GCS bucket that you created\n",
    "    !gsutil cp -n $FOLDER_NAME/books/* $GCS_BUCKET_URI/books\n",
    "    \n",
    "    print(f\"\\n{GCS_BUCKET_NAME} already exists. Contents:\\n\")\n",
    "    \n",
    "    # Verify that all Books 1 to 7 are uploaded to the GCS bucket (8 files in total, 2 for Part 1)\n",
    "    !gsutil ls $GCS_BUCKET_URI\n",
    "    \n",
    "def gcs_file(blob_name):\n",
    "    return bucket.blob(blob_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52e3ebf2-c792-4d04-9379-6dc1097bf0f3",
   "metadata": {},
   "source": [
    "# Part 1: Document Processing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a94b7aca-8fe8-4e91-ac9e-d578e16a2dbe",
   "metadata": {},
   "source": [
    "## Using Open Source Method to Process PDF Documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee4a9200-ac8a-4b6e-8b13-f04f7b9a7b78",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pdfreader import PDFDocument, SimplePDFViewer\n",
    "from typing import Optional\n",
    "from google.api_core.client_options import ClientOptions\n",
    "from google.cloud import discoveryengine\n",
    "\n",
    "# Load the PDF document\n",
    "pdf_file = gcs_file(\"books/Book1_HarryPotter_and_the_Sorcerers_Stone_pg15.pdf\")\n",
    "\n",
    "fd = pdf_file.open(\"rb\")\n",
    "doc = PDFDocument(fd)\n",
    "\n",
    "from io import BytesIO\n",
    "with pdf_file.open(\"rb\") as f:\n",
    "    stream = BytesIO(f.read())\n",
    "doc2 = PDFDocument(stream)\n",
    "\n",
    "page_one = next(doc.pages())\n",
    "\n",
    "all_pages = [p for p in doc.pages()]\n",
    "print(f\"Number of pages: {len(all_pages)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b0b3436-d991-4f8b-90d5-faddc8ec4a16",
   "metadata": {},
   "source": [
    "## Using Document AI to Process PDF Documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32db6ba2-21d4-4f38-b698-93db4eff6a17",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# !mkdir ./books\n",
    "# !mkdir ./matchingengine\n",
    "# !mkdir ./matchingengine/embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2b31f35-e7c8-4aac-b69b-ef11e57f6d23",
   "metadata": {},
   "source": [
    "### Define helper functions for processors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6767d615-1f74-4503-af3d-e152541523e1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Function to create Document AI Processor\n",
    "def create_processor(project_id, location, processor_display_name, processor_type):\n",
    "    # You must set the api_endpoint if you use a location other than 'us'.\n",
    "    opts = ClientOptions(api_endpoint=f\"{location}-documentai.googleapis.com\")\n",
    "\n",
    "    client = documentai.DocumentProcessorServiceClient(client_options=opts)\n",
    "\n",
    "    # The full resource name of the location\n",
    "    # e.g.: projects/project_id/locations/location\n",
    "    parent = client.common_location_path(project_id, location)\n",
    "\n",
    "    # Create a processor\n",
    "    processor = client.create_processor(\n",
    "        parent=parent,\n",
    "        processor=documentai.Processor(\n",
    "            display_name=processor_display_name, type_=processor_type\n",
    "        ),\n",
    "    )\n",
    "    \n",
    "    processor_id = processor.name.split('/')[-1]\n",
    "\n",
    "    # Print the processor information\n",
    "    print(f\"Processor Name: {processor.name}\")\n",
    "    print(f\"Processor Display Name: {processor.display_name}\")\n",
    "    print(f\"Processor ID: {processor_id}\")\n",
    "    print(f\"Processor Type: {processor.type_}\")\n",
    "    \n",
    "    \n",
    "    return processor, processor_id\n",
    "\n",
    "#Function to retrieve list of existing processors\n",
    "def list_processors(project_id: str, location: str) -> None:\n",
    "    processorList=[]\n",
    "    \n",
    "    # You must set the api_endpoint if you use a location other than 'us'.\n",
    "    opts = ClientOptions(api_endpoint=f\"{location}-documentai.googleapis.com\")\n",
    "\n",
    "    client = documentai.DocumentProcessorServiceClient(client_options=opts)\n",
    "\n",
    "    # The full resource name of the location\n",
    "    # e.g.: projects/project_id/locations/location\n",
    "    parent = client.common_location_path(PROJECT_ID, location)\n",
    "\n",
    "    # Make ListProcessors request\n",
    "    processor_list = client.list_processors(parent=parent)\n",
    "\n",
    "    # Print the processor information\n",
    "    for processor in processor_list:\n",
    "        # print(f\"Processor Name: {processor.name}\")\n",
    "        # print(f\"Processor Display Name: {processor.display_name}\")\n",
    "        # print(f\"Processor Type: {processor.type_}\")\n",
    "        processorList.append(processor)\n",
    "        \n",
    "    return processorList"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbfe4c03-dc3d-4c4c-9a5c-d001f950170a",
   "metadata": {},
   "source": [
    "### Import Document AI libraries and set variables\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "064db628-8878-4432-9593-1e4e06af4c1a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from google.api_core.client_options import ClientOptions\n",
    "from google.cloud import documentai\n",
    "\n",
    "GCP_PROJECT = PROJECT_ID #'cloud-llm-preview1'\n",
    "GCP_REGION='asia-southeast1'\n",
    "\n",
    "# Variables for Document AI OCR Processor\n",
    "PROCESSOR_DISPLAY_NAME = UNIQUE_PREFIX + '-ocr-processor' # Must be unique per project, e.g.: 'My Processor'\n",
    "PROCESSOR_TYPE = 'OCR_PROCESSOR' # Use fetch_processor_types to get available processor types"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dac09d5-e492-4d13-a247-d2d9ba4624ca",
   "metadata": {},
   "source": [
    "### Create Document AI Document OCR Processor\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6d44e0a-6322-4b04-8a49-d41fe27cda9c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "LOCATION=\"us\"\n",
    "\n",
    "processorList=list_processors(PROJECT_ID,LOCATION)\n",
    "\n",
    "def createUniqueProcessor():\n",
    "    if len(processorList)==0:\n",
    "        PROCESSOR, PROCESSOR_ID = create_processor(PROJECT_ID, LOCATION,PROCESSOR_DISPLAY_NAME, PROCESSOR_TYPE)\n",
    "        return PROCESSOR, PROCESSOR_ID\n",
    "    else:\n",
    "        for processor in processorList:\n",
    "            if PROCESSOR_DISPLAY_NAME==processor.display_name:\n",
    "                PROCESSOR_ID=processor.name.split('/')[-1]\n",
    "                PROCESSOR=processor\n",
    "                return PROCESSOR, PROCESSOR_ID\n",
    "            else:\n",
    "                try:\n",
    "                    PROCESSOR, PROCESSOR_ID = create_processor(PROJECT_ID, LOCATION,PROCESSOR_DISPLAY_NAME, PROCESSOR_TYPE)\n",
    "                    return PROCESSOR, PROCESSOR_ID\n",
    "                except:\n",
    "                    continue\n",
    "                \n",
    "\n",
    "PROCESSOR, PROCESSOR_ID = createUniqueProcessor()\n",
    "\n",
    "\n",
    "print(f\"Processor {PROCESSOR_ID} already exists.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c41daf7-2f20-4034-9087-2208b82d1aa9",
   "metadata": {},
   "source": [
    "### Processing a Single PDF Document using DocAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f51a8576-a992-4cc3-8216-e88a9276b4db",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from google.api_core.client_options import ClientOptions\n",
    "from google.cloud import documentai\n",
    "\n",
    "\n",
    "# PROJECT_ID = \"YOUR_PROJECT_ID\"\n",
    "GCP_PROJECT= PROJECT_ID #'cloud-llm-preview1'\n",
    "LOCATION = \"us\"  # Format is 'us' or 'eu'\n",
    "PROCESSOR_ID = PROCESSOR_ID  # Create processor in Cloud Console\n",
    "GCP_REGION=\"asia-southeast1\"\n",
    "\n",
    "# The local file in your current working directory\n",
    "FILE_PATH = file_path = gcs_file(\"books/Book1_HarryPotter_and_the_Sorcerers_Stone_pg15.pdf\")\n",
    "# Refer to https://cloud.google.com/document-ai/docs/file-types\n",
    "# for supported file types\n",
    "MIME_TYPE = mime_type = \"application/pdf\"\n",
    "\n",
    "# Instantiates a client\n",
    "docai_client = documentai.DocumentProcessorServiceClient(\n",
    "    client_options=ClientOptions(api_endpoint=f\"{LOCATION}-documentai.googleapis.com\")\n",
    ")\n",
    "\n",
    "# The full resource name of the processor, e.g.:\n",
    "# projects/project-id/locations/location/processor/processor-id\n",
    "# You must create new processors in the Cloud Console first\n",
    "RESOURCE_NAME = docai_client.processor_path(PROJECT_ID, LOCATION, PROCESSOR_ID)\n",
    "\n",
    "# Read the file into memory\n",
    "with FILE_PATH.open(\"rb\") as image:\n",
    "    image_content = image.read()\n",
    "\n",
    "# Load Binary Data into Document AI RawDocument Object\n",
    "raw_document = documentai.RawDocument(content=image_content, mime_type=MIME_TYPE)\n",
    "\n",
    "# Configure the process request\n",
    "request = documentai.ProcessRequest(name=RESOURCE_NAME, raw_document=raw_document)\n",
    "\n",
    "# Use the Document AI client to process the sample form\n",
    "result = docai_client.process_document(request=request)\n",
    "\n",
    "document_object = result.document\n",
    "print(\"Document processing complete.\")\n",
    "print(f\"Text: {document_object.text}\")\n",
    "\n",
    "page_text =document_object.text"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbf82e79-2d0f-45c3-9d09-24f071bc2a1b",
   "metadata": {},
   "source": [
    "### Using batch mode for processing multiple documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6b3aacd-e5f5-41a5-af02-2b1ec5218738",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import re\n",
    "from typing import Optional\n",
    "\n",
    "from google.api_core.client_options import ClientOptions\n",
    "from google.api_core.exceptions import InternalServerError\n",
    "from google.api_core.exceptions import RetryError\n",
    "from google.cloud import documentai  # type: ignore\n",
    "from google.cloud import storage\n",
    "import time\n",
    "\n",
    "#Creating batch processing function\n",
    "def batch_process_documents(\n",
    "    project_id: str,\n",
    "    location: str,\n",
    "    processor_id: str,\n",
    "    gcs_output_uri: str,\n",
    "    processor_version_id: Optional[str] = None,\n",
    "    gcs_input_uri: Optional[str] = None,\n",
    "    input_mime_type: Optional[str] = None,\n",
    "    gcs_input_prefix: Optional[str] = None,\n",
    "    field_mask: \"pages.pageNumber\" = 221,\n",
    "    timeout: int = 4000000000,\n",
    ") -> None:\n",
    "    # You must set the `api_endpoint` if you use a location other than \"us\".\n",
    "    opts = ClientOptions(api_endpoint=f\"{location}-documentai.googleapis.com\")\n",
    "\n",
    "    client = documentai.DocumentProcessorServiceClient(client_options=opts)\n",
    "\n",
    "    if gcs_input_uri:\n",
    "        # Specify specific GCS URIs to process individual documents\n",
    "        gcs_document = documentai.GcsDocument(\n",
    "            gcs_uri=gcs_input_uri, mime_type=input_mime_type\n",
    "        )\n",
    "        # Load GCS Input URI into a List of document files\n",
    "        gcs_documents = documentai.GcsDocuments(documents=[gcs_document])\n",
    "        input_config = documentai.BatchDocumentsInputConfig(gcs_documents=gcs_documents)\n",
    "    else:\n",
    "        # Specify a GCS URI Prefix to process an entire directory\n",
    "        gcs_prefix = documentai.GcsPrefix(gcs_uri_prefix=gcs_input_prefix)\n",
    "        input_config = documentai.BatchDocumentsInputConfig(gcs_prefix=gcs_prefix)\n",
    "\n",
    "    # Cloud Storage URI for the Output Directory\n",
    "    gcs_output_config = documentai.DocumentOutputConfig.GcsOutputConfig(\n",
    "        gcs_uri=gcs_output_uri, field_mask=field_mask\n",
    "    )\n",
    "\n",
    "    # Where to write results\n",
    "    output_config = documentai.DocumentOutputConfig(gcs_output_config=gcs_output_config)\n",
    "\n",
    "    if processor_version_id:\n",
    "        # The full resource name of the processor version, e.g.:\n",
    "        # projects/{project_id}/locations/{location}/processors/{processor_id}/processorVersions/{processor_version_id}\n",
    "        name = client.processor_version_path(\n",
    "            project_id, location, processor_id, processor_version_id\n",
    "        )\n",
    "    else:\n",
    "        # The full resource name of the processor, e.g.:\n",
    "        # projects/{project_id}/locations/{location}/processors/{processor_id}\n",
    "        name = client.processor_path(project_id, location, processor_id)\n",
    "\n",
    "    request = documentai.BatchProcessRequest(\n",
    "        name=name,\n",
    "        input_documents=input_config,\n",
    "        document_output_config=output_config,\n",
    "    )\n",
    "\n",
    "    # BatchProcess returns a Long Running Operation (LRO)\n",
    "    operation = client.batch_process_documents(request)\n",
    "\n",
    "    # Continually polls the operation until it is complete.\n",
    "    # This could take some time for larger files\n",
    "    # Format: projects/{project_id}/locations/{location}/operations/{operation_id}\n",
    "    try:\n",
    "        print(f\"Waiting for operation {operation.operation.name} to complete...\")\n",
    "        operation.result(timeout=timeout)\n",
    "    # Catch exception when operation doesn't finish before timeout\n",
    "    except (RetryError, InternalServerError) as e:\n",
    "        print(e.message)\n",
    "\n",
    "    # Once the operation is complete,\n",
    "    # get output document information from operation metadata\n",
    "    metadata = documentai.BatchProcessMetadata(operation.metadata)\n",
    "\n",
    "    if metadata.state != documentai.BatchProcessMetadata.State.SUCCEEDED:\n",
    "        raise ValueError(f\"Batch Process Failed: {metadata.state_message}\")\n",
    "\n",
    "    storage_client = storage.Client()\n",
    "\n",
    "    print(\"Output files:\")\n",
    "    output_document = []\n",
    "    # One process per Input Document\n",
    "    for process in list(metadata.individual_process_statuses):\n",
    "        # output_gcs_destination format: gs://BUCKET/PREFIX/OPERATION_NUMBER/INPUT_FILE_NUMBER/\n",
    "        # The Cloud Storage API requires the bucket name and URI prefix separately\n",
    "        matches = re.match(r\"gs://(.*?)/(.*)\", process.output_gcs_destination)\n",
    "        if not matches:\n",
    "            print(\n",
    "                \"Could not parse output GCS destination:\",\n",
    "                process.output_gcs_destination,\n",
    "            )\n",
    "            continue\n",
    "\n",
    "        output_bucket, output_prefix = matches.groups()\n",
    "\n",
    "        # Get List of Document Objects from the Output Bucket\n",
    "        output_blobs = storage_client.list_blobs(output_bucket, prefix=output_prefix)\n",
    "        \n",
    "        # Document AI may output multiple JSON files per source file\n",
    "        for blob in output_blobs:\n",
    "            # Document AI should only output JSON files to GCS\n",
    "            if blob.content_type != \"application/json\":\n",
    "                print(\n",
    "                    f\"Skipping non-supported file: {blob.name} - Mimetype: {blob.content_type}\"\n",
    "                )\n",
    "                continue\n",
    "\n",
    "            # Download JSON File as bytes object and convert to Document Object\n",
    "            print(f\"Fetching {blob.name}\")\n",
    "            document = documentai.Document.from_json(\n",
    "                blob.download_as_bytes(), ignore_unknown_fields=True\n",
    "            )\n",
    "\n",
    "            # For a full list of Document object attributes, please reference this page:\n",
    "            # https://cloud.google.com/python/docs/reference/documentai/latest/google.cloud.documentai_v1.types.Document\n",
    "            \n",
    "            # Read the text recognition output from the processor \n",
    "            print(\"The document contains the following text:\")\n",
    "            print(document.text)\n",
    "            output_document.append(document.text)\n",
    "    return(\"\".join(output_document))\n",
    "\n",
    "def save_text_to_file(text, filename):\n",
    "    pattern = r\".*/([^/.]+)\\.pdf\"\n",
    "\n",
    "    # Extract the filename\n",
    "    match = re.search(pattern, filename)\n",
    "    if match:\n",
    "        filename = match.group(1)\n",
    "        print(filename + \" has been processed successfully.\\n\")  # Output ex: Book4_The_Goblet_of_Fire\n",
    "    else:\n",
    "        print(\"No match found\")\n",
    "    \n",
    "    file = gcs_file(\"results/\" + filename + \".txt\")\n",
    "    \n",
    "    with file.open('w', encoding='utf-8') as f:\n",
    "        f.write(text)    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "707c20f5-b4d6-4114-9ed6-689df0232389",
   "metadata": {},
   "source": [
    "### Setting Variables & Triggering the Batch Processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3839fca8-391d-44a3-ade1-28396ccd331a",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# PROJECT_ID = \"YOUR_PROJECT_ID\"\n",
    "GCP_PROJECT= PROJECT_ID #'cloud-llm-preview1'\n",
    "LOCATION = location = \"us\"  # Format is 'us' or 'eu'\n",
    "project_id=PROJECT_ID\n",
    "processor_id=PROCESSOR_ID\n",
    "\n",
    "# The local file in your current working directory\n",
    "FILE_PATH = f\"{GCS_BUCKET_URI}/books/Book1_The_Sorcerers_Stone.pdf\"\n",
    "# Refer to https://cloud.google.com/document-ai/docs/file-types\n",
    "# for supported file types\n",
    "\n",
    "# TODO(developer): Uncomment these variables before running the sample.\n",
    "gcs_output_uri = f\"{GCS_BUCKET_URI}/\" # Must end with a trailing slash `/`. Format: gs://bucket/directory/subdirectory/\n",
    "# processor_version_id = \"\" # Optional. Example: pretrained-ocr-v1.0-2020-09-23\n",
    "\n",
    "# TODO(developer): You must specify either `gcs_input_uri` and `mime_type` or `gcs_input_prefix`\n",
    "gcs_input_uri = f\"{GCS_BUCKET_URI}/books/Book1_The_Sorcerers_Stone.pdf\" # Format: gs://bucket/directory/file.pdf\n",
    "MIME_TYPE = input_mime_type = \"application/pdf\"\n",
    "\n",
    "gcs_input_prefix = f\"{GCS_BUCKET_URI}/matchingengine/\" # Format: gs://bucket/directory/\n",
    "field_mask = \"text,entities,pages.pageNumber\"  # Optional. The fields to return in the Document object.\n",
    "timeout = 400000\n",
    "\n",
    "book_list = [f\"{GCS_BUCKET_URI}/books/Book1_The_Sorcerers_Stone.pdf\",]\n",
    "             # f\"{GCS_BUCKET_URI}/books/Book2_The_Chamber_of_Secrets.pdf\",\n",
    "             # f\"{GCS_BUCKET_URI}/books/Book3_The_Prisoner_of_Azkaban.pdf\",\n",
    "             # f\"{GCS_BUCKET_URI}/books/Book4_The_Goblet_of_Fire.pdf\",\n",
    "             # f\"{GCS_BUCKET_URI}/books/Book5_The_Order_of_the_Phoenix.pdf\",\n",
    "             # f\"{GCS_BUCKET_URI}/books/Book6_The_HalfBlood_Prince.pdf\",\n",
    "             # f\"{GCS_BUCKET_URI}/books/Book7_The_Deathly_Hallows.pdf\",]\n",
    "try:\n",
    "    for i in range(0,len(book_list)): \n",
    "        gcs_input_uri = book_list[i]\n",
    "        print(gcs_input_uri + \":\\n\")\n",
    "\n",
    "        page_text_batch = batch_process_documents(   project_id,\n",
    "            location,\n",
    "            processor_id,\n",
    "            gcs_output_uri,\n",
    "            None,\n",
    "            gcs_input_uri,\n",
    "            input_mime_type,\n",
    "            gcs_input_prefix,\n",
    "            field_mask,\n",
    "            timeout)\n",
    "\n",
    "        # Example usage: my_text = \"This is the text from your OCR process.\"\n",
    "\n",
    "        save_text_to_file(page_text_batch, gcs_input_uri)\n",
    "\n",
    "\n",
    "        time.sleep(60)\n",
    "except:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b71d65a-e393-4af9-8716-b4c836311d1b",
   "metadata": {},
   "source": [
    "#### Here is the full documenation for DocAI: https://cloud.google.com/document-ai/docs/samples/documentai-batch-process-document?hl=en\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff99b44d-c70f-4ca4-8dd2-d49782243627",
   "metadata": {},
   "source": [
    "## Part 2: Embeddings API"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "159b7a38-b79d-4d6f-a8e1-4a4042cff3cd",
   "metadata": {},
   "source": [
    "### Initialising Vertex AI & Setting Up Embeddings API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80e1de11-3627-4081-832e-b8c40c1e0af3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# init the vertexai package\n",
    "import vertexai\n",
    "LOCATION=\"us-central1\"\n",
    "\n",
    "vertexai.init(project=PROJECT_ID, location=LOCATION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7ffde96-8b94-468e-beff-bee6ae841505",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load the text embeddings model\n",
    "from vertexai.preview.language_models import TextEmbeddingModel \n",
    "embedding_model =\"textembedding-gecko@003\"\n",
    "model = TextEmbeddingModel.from_pretrained(embedding_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a8892fb-e899-4cfc-8c4a-1489bba7232d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import tqdm  # to show a progress bar\n",
    "import os\n",
    "# get embeddings for a list of texts\n",
    "BATCH_SIZE = 5\n",
    "\n",
    "\n",
    "def get_embeddings_wrapper(texts):\n",
    "    embs = []\n",
    "    for i in tqdm.tqdm(range(0, len(texts), BATCH_SIZE)):\n",
    "        time.sleep(1)  # to avoid the quota error\n",
    "        result = model.get_embeddings(texts[i : i + BATCH_SIZE])\n",
    "        embs = embs + [e.values for e in result]\n",
    "    return embs[0]\n",
    "\n",
    "def text_embedding(text) -> list:\n",
    "    \"\"\"Text embedding with a Large Language Model.\"\"\"\n",
    "    model = TextEmbeddingModel.from_pretrained(embedding_model)\n",
    "    embeddings = model.get_embeddings(text)\n",
    "    for embedding in embeddings:\n",
    "        vector = embedding.values\n",
    "        print(f\"Length of Embedding Vector: {len(vector)}\")\n",
    "    return vector"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa1899b9-9285-4014-8c35-46d65bcd840b",
   "metadata": {},
   "source": [
    "### Convert text into embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fba116f0-ce74-479d-8d88-da7dbd96d97a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "trail_text = [page_text]\n",
    "\n",
    "text_vectors = text_embedding(trail_text)\n",
    "\n",
    "#Example of what the embeddings look like\n",
    "for i in range(0,15):\n",
    "    print(text_vectors[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74c0d878-36b4-401f-b620-40290ee80ede",
   "metadata": {},
   "source": [
    "### Split text into lines (or sentences/words separated by symbols)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70c5c11a-6ba1-4a72-8d12-4fab7f310be3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter(\n",
    "    # Set a really small chunk size, just to show.\n",
    "    chunk_size=100,\n",
    "    chunk_overlap=20,\n",
    "    separators= [\"/,\", \"##\", \">\", \"-\"],#'\\n\\n', '\\n'],\n",
    "    length_function=len,\n",
    "    is_separator_regex=False,\n",
    ")\n",
    "\n",
    "splitted_texts = text_splitter.create_documents([page_text])\n",
    "print(\"one \\n\",splitted_texts[0])\n",
    "print(\"two \\n\",splitted_texts[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dcdc58e-f908-4200-9c68-3d80bd3c1356",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Display text that has been split in a DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e96cdc0-f358-4903-b7ee-4fa30cfd57da",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "splitted_texts_list = text_splitter.split_text(page_text)#[:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99c98cc9-53fe-4244-8546-6f5c8ca716ec",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "# Create a DataFrame from the splitted texts\n",
    "df = pd.DataFrame({'splitted_texts': splitted_texts_list})\n",
    "\n",
    "# Add a row number column\n",
    "df['id'] = df.index + 1\n",
    "\n",
    "# Print the DataFrame\n",
    "print(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49f98953-5f13-482b-aca2-52bd6b545eb2",
   "metadata": {},
   "source": [
    "### Display text that has been split in a List"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efdb51db-8dc1-4293-b897-d36082df8c70",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "list(df.splitted_texts)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f899b32d-02b0-4ad3-8382-f29298c503d4",
   "metadata": {},
   "source": [
    "# Part 3: Matching Engine"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f607703-fa6d-411f-867a-b6d68f5180f5",
   "metadata": {},
   "source": [
    "## Using Document AI Online Processing to process PDF document"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6db79b7a-0057-4c65-9a09-e4f8f1aeed44",
   "metadata": {},
   "source": [
    "### Declare DocAI Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c3e093d-2fc7-410e-ac4c-105eaa7cac69",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "LOCATION = location = \"us\"\n",
    "processor_version = 'rc'\n",
    "\n",
    "def online_process(\n",
    "    project_id: str,\n",
    "    location: str,\n",
    "    processor_id: str,\n",
    "    file_path: str,\n",
    "    mime_type: str,\n",
    ") -> documentai.Document:\n",
    "    \"\"\"\n",
    "    A function to process a document online using Google Document AI.\n",
    "    \"\"\"\n",
    "\n",
    "    # Define an options dictionary, which includes the API's URL. This is used to connect to Google's Document AI service\n",
    "    opts = {\"api_endpoint\": f\"{location}-documentai.googleapis.com\"}\n",
    "\n",
    "    # Create a Document AI client, think of it as our bridge for communicating with Google's services\n",
    "    documentai_client = documentai.DocumentProcessorServiceClient(client_options=opts)\n",
    "\n",
    "    # Generate the complete name of the processor\n",
    "    # You need to first create a processor in the Google Cloud console\n",
    "    resource_name = documentai_client.processor_path(project_id, location, processor_id)\n",
    "\n",
    "    # Read in the document you want to analyze (like an image or PDF), and store it in the variable image_content\n",
    "    with file_path.open(\"rb\") as image:\n",
    "        image_content = image.read()\n",
    "\n",
    "        # Convert the read document into a format that Google Document AI can understand, i.e., a RawDocument object\n",
    "        raw_document = documentai.RawDocument(\n",
    "            content=image_content, mime_type=mime_type\n",
    "        )\n",
    "        # Create a request, which includes the name of the processor and the document we want to analyze\n",
    "        request = documentai.ProcessRequest(\n",
    "            name=resource_name, raw_document=raw_document\n",
    "        )\n",
    "        # Send our request and receive the analysis results\n",
    "        result = documentai_client.process_document(request=request)\n",
    "        \n",
    "        print(\"Document processing complete.\")\n",
    "        # print(f\"Text: {document_object.text}\")\n",
    "        \n",
    "        # Return this analysis result\n",
    "        return result.document\n",
    "\n",
    "#Remove unwanted characters    \n",
    "def trim_text(text: str): \n",
    "    \"\"\" Removes spaces and newline characters. \"\"\" \n",
    "    return text.strip().replace(\"\\n\", \" \")\n",
    "\n",
    "def process_document(\n",
    "    project_id: str,\n",
    "    location: str,\n",
    "    processor_id: str,\n",
    "    processor_version: str,\n",
    "    file_path: str,\n",
    "    mime_type: str,\n",
    "    process_options: Optional[documentai.ProcessOptions] = None,\n",
    ") -> documentai.Document:\n",
    "    # You must set the `api_endpoint` if you use a location other than \"us\".\n",
    "    client = documentai.DocumentProcessorServiceClient(\n",
    "        client_options=ClientOptions(\n",
    "            api_endpoint=f\"{location}-documentai.googleapis.com\"\n",
    "        )\n",
    "    )\n",
    "\n",
    "    # The full resource name of the processor version, e.g.:\n",
    "    # `projects/{project_id}/locations/{location}/processors/{processor_id}/processorVersions/{processor_version_id}`\n",
    "    # You must create a processor before running this sample.\n",
    "    name = client.processor_version_path(\n",
    "        project_id, location, processor_id, processor_version\n",
    "    )\n",
    "\n",
    "    # Read the file into memory\n",
    "    with open(file_path, \"rb\") as image:\n",
    "        image_content = image.read()\n",
    "\n",
    "    # Configure the process request\n",
    "    request = documentai.ProcessRequest(\n",
    "        name=name,\n",
    "        raw_document=documentai.RawDocument(content=image_content, mime_type=mime_type),\n",
    "        # Only supported for Document OCR processor\n",
    "        process_options=process_options,\n",
    "    )\n",
    "\n",
    "    result = client.process_document(request=request)\n",
    "\n",
    "    # For a full list of `Document` object attributes, reference this page:\n",
    "    # https://cloud.google.com/document-ai/docs/reference/rest/v1/Document\n",
    "    return result.document"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf6ff9f9-3ef3-40da-8aa1-540bd0874b55",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Extra DocAI Functions for reference (Will not be used in this lab)\n",
    "#### Documentation: https://cloud.google.com/document-ai/docs/handle-response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e20031f-36ac-4bed-b4b9-f8722e026270",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# from typing import Optional, Sequence\n",
    "\n",
    "# from google.api_core.client_options import ClientOptions\n",
    "# from google.cloud import documentai\n",
    "\n",
    "# def print_page_dimensions(dimension: documentai.Document.Page.Dimension) -> None:\n",
    "#     print(f\"    Width: {str(dimension.width)}\")\n",
    "#     print(f\"    Height: {str(dimension.height)}\")\n",
    "\n",
    "# def print_detected_langauges(\n",
    "#     detected_languages: Sequence[documentai.Document.Page.DetectedLanguage],\n",
    "# ) -> None:\n",
    "#     print(\"    Detected languages:\")\n",
    "#     for lang in detected_languages:\n",
    "#         print(f\"        {lang.language_code} ({lang.confidence:.1%} confidence)\")\n",
    "\n",
    "\n",
    "# def print_blocks(blocks: Sequence[documentai.Document.Page.Block], text: str) -> None:\n",
    "#     print(f\"    {len(blocks)} blocks detected:\")\n",
    "#     first_block_text = layout_to_text(blocks[0].layout, text)\n",
    "#     print(f\"        First text block: {repr(first_block_text)}\")\n",
    "#     last_block_text = layout_to_text(blocks[-1].layout, text)\n",
    "#     print(f\"        Last text block: {repr(last_block_text)}\")\n",
    "\n",
    "\n",
    "# def print_paragraphs(\n",
    "#     paragraphs: Sequence[documentai.Document.Page.Paragraph], text: str\n",
    "# ) -> None:\n",
    "#     print(f\"    {len(paragraphs)} paragraphs detected:\")\n",
    "#     first_paragraph_text = layout_to_text(paragraphs[0].layout, text)\n",
    "#     print(f\"        First paragraph text: {repr(first_paragraph_text)}\")\n",
    "\n",
    "#     last_paragraph_text = layout_to_text(paragraphs[-1].layout, text)\n",
    "#     print(f\"        Last paragraph text: {repr(last_paragraph_text)}\")\n",
    "\n",
    "\n",
    "# def print_lines(lines: Sequence[documentai.Document.Page.Line], text: str) -> None:\n",
    "#     print(f\"    {len(lines)} lines detected:\")\n",
    "#     first_line_text = layout_to_text(lines[0].layout, text)\n",
    "#     print(f\"        First line text: {repr(first_line_text)}\")\n",
    "#     last_line_text = layout_to_text(lines[-1].layout, text)\n",
    "#     print(f\"        Last line text: {repr(last_line_text)}\")\n",
    "\n",
    "\n",
    "# def print_tokens(tokens: Sequence[documentai.Document.Page.Token], text: str) -> None:\n",
    "#     print(f\"    {len(tokens)} tokens detected:\")\n",
    "#     first_token_text = layout_to_text(tokens[0].layout, text)\n",
    "#     first_token_break_type = tokens[0].detected_break.type_.name\n",
    "#     print(f\"        First token text: {repr(first_token_text)}\")\n",
    "#     print(f\"        First token break type: {repr(first_token_break_type)}\")\n",
    "#     if tokens[0].style_info:\n",
    "#         print_style_info(tokens[0].style_info)\n",
    "\n",
    "#     last_token_text = layout_to_text(tokens[-1].layout, text)\n",
    "#     last_token_break_type = tokens[-1].detected_break.type_.name\n",
    "#     print(f\"        Last token text: {repr(last_token_text)}\")\n",
    "#     print(f\"        Last token break type: {repr(last_token_break_type)}\")\n",
    "#     if tokens[-1].style_info:\n",
    "#         print_style_info(tokens[-1].style_info)\n",
    "\n",
    "\n",
    "# def print_symbols(\n",
    "#     symbols: Sequence[documentai.Document.Page.Symbol], text: str\n",
    "# ) -> None:\n",
    "#     print(f\"    {len(symbols)} symbols detected:\")\n",
    "#     first_symbol_text = layout_to_text(symbols[0].layout, text)\n",
    "#     print(f\"        First symbol text: {repr(first_symbol_text)}\")\n",
    "#     last_symbol_text = layout_to_text(symbols[-1].layout, text)\n",
    "#     print(f\"        Last symbol text: {repr(last_symbol_text)}\")\n",
    "\n",
    "\n",
    "# def print_image_quality_scores(\n",
    "#     image_quality_scores: documentai.Document.Page.ImageQualityScores,\n",
    "# ) -> None:\n",
    "#     print(f\"    Quality score: {image_quality_scores.quality_score:.1%}\")\n",
    "#     print(\"    Detected defects:\")\n",
    "\n",
    "#     for detected_defect in image_quality_scores.detected_defects:\n",
    "#         print(f\"        {detected_defect.type_}: {detected_defect.confidence:.1%}\")\n",
    "\n",
    "\n",
    "# def print_style_info(style_info: documentai.Document.Page.Token.StyleInfo) -> None:\n",
    "#     \"\"\"\n",
    "#     Only supported in version `pretrained-ocr-v2.0-2023-06-02`\n",
    "#     \"\"\"\n",
    "#     print(f\"           Font Size: {style_info.font_size}pt\")\n",
    "#     print(f\"           Font Type: {style_info.font_type}\")\n",
    "#     print(f\"           Bold: {style_info.bold}\")\n",
    "#     print(f\"           Italic: {style_info.italic}\")\n",
    "#     print(f\"           Underlined: {style_info.underlined}\")\n",
    "#     print(f\"           Handwritten: {style_info.handwritten}\")\n",
    "#     print(\n",
    "#         f\"           Text Color (RGBa): {style_info.text_color.red}, {style_info.text_color.green}, {style_info.text_color.blue}, {style_info.text_color.alpha}\"\n",
    "#     )\n",
    "\n",
    "\n",
    "# def print_visual_elements(\n",
    "#     visual_elements: Sequence[documentai.Document.Page.VisualElement], text: str\n",
    "# ) -> None:\n",
    "#     \"\"\"\n",
    "#     Only supported in version `pretrained-ocr-v2.0-2023-06-02`\n",
    "#     \"\"\"\n",
    "#     checkboxes = [x for x in visual_elements if \"checkbox\" in x.type]\n",
    "#     math_symbols = [x for x in visual_elements if x.type == \"math_formula\"]\n",
    "\n",
    "#     if checkboxes:\n",
    "#         print(f\"    {len(checkboxes)} checkboxes detected:\")\n",
    "#         print(f\"        First checkbox: {repr(checkboxes[0].type)}\")\n",
    "#         print(f\"        Last checkbox: {repr(checkboxes[-1].type)}\")\n",
    "\n",
    "#     if math_symbols:\n",
    "#         print(f\"    {len(math_symbols)} math symbols detected:\")\n",
    "#         first_math_symbol_text = layout_to_text(math_symbols[0].layout, text)\n",
    "#         print(f\"        First math symbol: {repr(first_math_symbol_text)}\")\n",
    "\n",
    "# def layout_to_text(layout: documentai.Document.Page.Layout, text: str) -> str:\n",
    "#     \"\"\"\n",
    "#     Document AI identifies text in different parts of the document by their\n",
    "#     offsets in the entirety of the document\"s text. This function converts\n",
    "#     offsets to a string.\n",
    "#     \"\"\"\n",
    "#     # If a text segment spans several lines, it will\n",
    "#     # be stored in different text segments.\n",
    "#     return \"\".join(\n",
    "#         text[int(segment.start_index) : int(segment.end_index)]\n",
    "#         for segment in layout.text_anchor.text_segments\n",
    "#     )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b69e857e-27f7-4aaa-8ce5-7c594f7be5dd",
   "metadata": {},
   "source": [
    "### Trigger DocAI Online Document Processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c56eb52-6050-46b7-899d-e075b7db4bf5",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "document = online_process(\n",
    "    project_id=project_id,\n",
    "    location=location,\n",
    "    processor_id=processor_id,\n",
    "    file_path=file_path,\n",
    "    mime_type=mime_type,\n",
    ")\n",
    "\n",
    "names = []\n",
    "name_confidence = []\n",
    "values = []\n",
    "value_confidence = []"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5d75a27-a23e-4b73-b7ef-5d094161ec63",
   "metadata": {},
   "source": [
    "## Using langchain for Doc AI on GCP"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77232c44-5666-4a4d-9e2c-ed0fc2d888ff",
   "metadata": {},
   "source": [
    "### Configure a DocAIParser "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b548093-a8ec-4e00-81c3-efadf6bc2f57",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# from langchain_community.document_loaders.blob_loaders import Blob\n",
    "# from langchain_community.document_loaders.parsers import DocAIParser\n",
    "\n",
    "# GCS_OUTPUT_PATH= GCS_BUCKET_URI \n",
    "# PROCESSOR_NAME = f'projects/{PROJECT_NUMBER}/locations/us/processors/{PROCESSOR_ID}'\n",
    "\n",
    "\n",
    "# parser = DocAIParser(\n",
    "#     location=\"us\", processor_name=PROCESSOR_NAME, gcs_output_path=GCS_OUTPUT_PATH\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25423d53-c063-4da3-a717-d41f54f16194",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def combine_text_files(text_files):  \n",
    "    combined_text = []\n",
    "    for file in text_files:\n",
    "        blob = gcs_file(file)\n",
    "        text = blob.download_as_text()\n",
    "        combined_text.append(text)\n",
    "    return \"\\n\\n\".join(combined_text)\n",
    "\n",
    "blob_names = [blob.name for blob in bucket.list_blobs(prefix = \"results/\")]\n",
    "\n",
    "text = combine_text_files(blob_names)\n",
    "page_contents = text_splitter.split_text(text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77a09418-070d-435e-8e0c-f666bb478845",
   "metadata": {},
   "source": [
    "### Display processed pages in a DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ffa6fb9-aaca-4319-a050-4e3b9141e8ae",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "# Create a DataFrame from the splitted texts\n",
    "df = pd.DataFrame({'pagewise_texts': page_contents})\n",
    "\n",
    "# Add a row number column\n",
    "df['page_id'] = df.index + 1\n",
    "\n",
    "# Print the DataFrame\n",
    "print(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02eb3ca3-d515-4c09-a403-da178630d916",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "chunk_size=100\n",
    "chunk_overlap=20\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter(\n",
    "    # Set a really small chunk size, just to show.\n",
    "    chunk_size=chunk_size,\n",
    "    chunk_overlap=chunk_overlap,\n",
    "    separators= [\"/,\", \"##\", \">\", \"We\"],#'\\n\\n', '\\n'],\n",
    "    length_function=len,\n",
    "    is_separator_regex=False,\n",
    ")\n",
    "\n",
    "def split_text_chunks(text, chunk_size):\n",
    "    return [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]\n",
    "\n",
    "splitted_texts = []\n",
    "\n",
    "for index, row in df.iterrows():\n",
    "    splitted_texts += text_splitter.create_documents(row['pagewise_texts'])\n",
    "    splitted_texts2 = text_splitter.split_text(row['pagewise_texts'])\n",
    "    \n",
    "    df['splitted_texts'] = df['pagewise_texts'].apply(lambda x: text_splitter.split_text(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "212bf2d4-3c94-450e-9773-70f2f5705780",
   "metadata": {},
   "source": [
    "### Display text that has been split in a DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "786d6e85-3164-4ec9-b926-02836c069504",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "890262cb-7c86-41ee-90d8-721adbb16ec9",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Display text that has been split in a Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68ce1979-d84e-4537-9ba5-a8c58159bcad",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_exploded = df.explode('splitted_texts')\n",
    "df_exploded.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcaeed0b-f91e-4430-b33b-6767bbdda899",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "len(splitted_texts)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba51058e-5469-4ba0-9774-7677c1a5f642",
   "metadata": {},
   "source": [
    "### Split Text into Chunks, and add it to the Table as a column"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "664ef2b0-ec7d-49bc-9f6e-880fea34027e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_exploded['splitted_texts_chunks'] = df_exploded['splitted_texts'].apply(lambda x: split_text_chunks(x,chunk_size))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c943c16e-101e-4046-89ba-2f7f2cb78bcf",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_exploded_2 = df_exploded.explode('splitted_texts_chunks')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "485c8def-e82b-40fe-a231-eff085234299",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_exploded_2.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddba6089-1c7f-492c-b00b-50bef95d1df2",
   "metadata": {},
   "source": [
    "### Index each row in the Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18959ac0-a3c6-473a-9df5-411150827975",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_exploded_2x = df_exploded_2.copy()\n",
    "\n",
    "df_exploded_2x = df_exploded_2x.reindex()\n",
    "df_exploded_2x = df_exploded_2x.reset_index()\n",
    "\n",
    "df_exploded_2x['id'] = df_exploded_2x.index\n",
    "df_exploded_2x['id'] = df_exploded_2x['id'].astype(str)\n",
    "\n",
    "df_exploded_2 = df_exploded_2x.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c4b21f6-5c28-4927-b7c4-23ac2d85ed1c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_exploded_2.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "902bce95-410d-49c2-84e5-d4e1cd9a287a",
   "metadata": {},
   "source": [
    "### Convert DataFrame (excluding Questions) into a CSV file\n",
    "Split the dataframe into chunks and upload the file to cloud storage\n",
    "(For the purposes of the workshop, we will be omitting this due to the time required to run this)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b700db6d-5ca8-4ede-a87a-73b17178a6f8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_exploded_2.to_csv('df_exploded_2.csv', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0404e68-95b8-4040-b739-25585d2705ac",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Import VertexAIEmbeddings and set the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47c5477a-bfc7-41ea-acf7-0a58e7fdc98b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain_google_vertexai import VertexAIEmbeddings\n",
    "\n",
    "embeddings = VertexAIEmbeddings(model_name=\"textembedding-gecko@003\")\n",
    "\n",
    "text = \"This is a test document.\"\n",
    "\n",
    "query_result = embeddings.embed_query(text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ff50e64-b554-4e11-8f54-9a626571f281",
   "metadata": {},
   "source": [
    "### Import Question Bank from BigQuery public dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6229134-f993-466f-9efe-1bd31966190e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# load the BQ Table into a Pandas Dataframe\n",
    "import pandas as pd\n",
    "from google.cloud import bigquery\n",
    "\n",
    "QUESTIONS_SIZE = df_exploded_2.shape[0] #1000\n",
    "\n",
    "bq_client = bigquery.Client(project=PROJECT_ID)\n",
    "QUERY_TEMPLATE = \"\"\"\n",
    "        SELECT distinct q.id, q.title\n",
    "        FROM (SELECT * FROM `bigquery-public-data.stackoverflow.posts_questions`\n",
    "        where Score > 0 ORDER BY View_Count desc) AS q\n",
    "        LIMIT {limit} ;\n",
    "        \"\"\"\n",
    "query = QUERY_TEMPLATE.format(limit=QUESTIONS_SIZE)\n",
    "query_job = bq_client.query(query)\n",
    "rows = query_job.result()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de493b20-2556-47ca-ab51-a2129a5f044b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "len(df_exploded_2.splitted_texts_chunks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a35d304b-44f9-46f8-a1ec-8262ee1f757c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_exploded_2.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8119bb3-d373-4b2d-999a-87a0a458dee2",
   "metadata": {},
   "source": [
    "### Convert text from questions into Embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f7120e4-0a59-428d-8a61-c6598456e8dd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import tqdm  # to show a progress bar\n",
    "from vertexai.language_models import TextEmbeddingModel\n",
    "\n",
    "# get embeddings for a list of texts\n",
    "BATCH_SIZE = 5\n",
    "\n",
    "\n",
    "model_ai=\"textembedding-gecko@003\"\n",
    "\n",
    "model = TextEmbeddingModel.from_pretrained(model_ai)\n",
    "\n",
    "def get_embeddings_wrapper(texts):\n",
    "    embs = []\n",
    "    for i in tqdm.tqdm(range(0, len(texts), BATCH_SIZE)):\n",
    "        time.sleep(1)  # to avoid the quota error\n",
    "        result = model.get_embeddings(texts[i : i + BATCH_SIZE])\n",
    "        embs = embs + [e.values for e in result]\n",
    "    return embs\n",
    "# The following code will get embedding for the question titles and add them as a new column embedding to the DataFrame. This will take a few minutes.\n",
    "\n",
    "# get embeddings for the question titles and add them as \"embedding\" column\n",
    "df = df_exploded_2.assign(embedding=get_embeddings_wrapper(list(df_exploded_2.splitted_texts_chunks)))\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0340f985-329b-4cea-a39c-4b465ae73c3c",
   "metadata": {},
   "source": [
    "### Convert Embeddings into a JSON File"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1c32146-5849-40f0-83a2-5ef24e734783",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# save id and embedding as a json file\n",
    "jsonl_string = df[[\"id\",'splitted_texts_chunks', \"embedding\"]].to_json(orient=\"records\", lines=True)\n",
    "with open(f\"{FOLDER_NAME}/questions_test.json\", \"w\") as f:\n",
    "    f.write(jsonl_string)\n",
    "\n",
    "# show the first few lines of the json file\n",
    "! head -n 3 questions_test.json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "015f09be-0070-4e7e-9273-3ae7808ea578",
   "metadata": {},
   "source": [
    "### Upload the JSON File to matching engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f33b95e-c997-4d52-b07c-be8b266da2a0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# generate an unique id for this session\n",
    "from datetime import datetime\n",
    "\n",
    "UID = UNIQUE_PREFIX\n",
    "\n",
    "BUCKET_URI_ME=f\"{GCS_BUCKET_URI}/matchingengine/embedding/\"\n",
    "LOCATION = 'asia-southeast1'\n",
    "\n",
    "from google.cloud import aiplatform\n",
    "\n",
    "aiplatform.init(project=PROJECT_ID, location=LOCATION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "201cd07b-6263-45cd-8849-5425d53ea89d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "! gsutil cp questions_test.json {BUCKET_URI_ME}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4679b3d5-4a07-490a-94e6-9de1a960f692",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "! gsutil ls {BUCKET_URI_ME}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca5d6483-a3cb-46b8-bdfa-aa7bb9bde54d",
   "metadata": {},
   "source": [
    "## Creating Matching Engine Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f71e82a-46d0-4bd7-9b7c-6a6c8daeb40b",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# create Index\n",
    "my_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(\n",
    "    display_name=f\"vs-feature-index-{UID}\",\n",
    "    contents_delta_uri=BUCKET_URI_ME,\n",
    "    dimensions=768,\n",
    "    approximate_neighbors_count=10,\n",
    "    project = PROJECT_ID\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d82ee84-4265-478d-a852-7cef90a1daf1",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Create Index Endpoint and deploy the Index\n",
    "To use the Index, you need to create an Index Endpoint. It works as a server instance accepting query requests for your Index.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16714e23-2f8c-4153-a9ba-67549975cc6f",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# create IndexEndpoint\n",
    "my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint.create(\n",
    "    display_name=f\"vs-feature-index-endpoint-{UID}\", public_endpoint_enabled=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b619bbe-ab6e-4b23-94bd-70a7c939cc16",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "DEPLOYED_INDEX_ID = f\"vs_feature_deployed_{UID}\"\n",
    "# deploy the Index to the Index Endpoint\n",
    "my_index_endpoint.deploy_index(index=my_index, deployed_index_id=DEPLOYED_INDEX_ID)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78a35991-3e36-4c01-bbf3-e6e03664a181",
   "metadata": {},
   "source": [
    "#### Go to you vertex AI console and check the index is CREATED successfully "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fdade17-fe17-41b4-8282-4da5f93e152d",
   "metadata": {},
   "source": [
    "### Querying a created index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82c8cf62-3686-4b8d-8e46-1d69d6106ebd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "# build dicts for product names and embs\n",
    "product_names = {}\n",
    "product_embs = {}\n",
    "product_text = {}\n",
    "with open(f\"{FOLDER_NAME}/questions_test.json\") as f:\n",
    "    for l in f.readlines():\n",
    "        p = json.loads(l)\n",
    "        id = p[\"id\"]\n",
    "        product_names[id] = p[\"id\"]\n",
    "        product_text[id] = p['splitted_texts_chunks']\n",
    "        product_embs[id] = p[\"embedding\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9af7516-1e0d-4ed5-a88a-4f15584e51e3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# get the embedding for ID 6523 \"cloudveil women's excursion short\"\n",
    "# you can also try with other IDs such as 12711, 18090, 19536 and 11863\n",
    "query_emb = product_embs[\"0\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8f50a05-df56-4ed1-9747-926a5a713575",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# run query\n",
    "response = my_index_endpoint.find_neighbors(\n",
    "    deployed_index_id=DEPLOYED_INDEX_ID, queries=[query_emb], num_neighbors=3\n",
    ")\n",
    "\n",
    "# show the results\n",
    "for idx, neighbor in enumerate(response[0]):\n",
    "    print(f\"{neighbor.distance:.2f} {product_names[neighbor.id]} {product_text[neighbor.id]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b31ec425-9063-4389-ac1e-f93b67c9f051",
   "metadata": {},
   "source": [
    "### Run Query\n",
    "Finally it's ready to use Vector Search. In the following code, it creates an embedding for a test question, and find similar question with the Vector Search."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72d134bd-905c-483c-9a4d-79e735da1e13",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import tqdm  # to show a progress bar\n",
    "\n",
    "# get embeddings for a list of texts\n",
    "BATCH_SIZE = 5\n",
    "\n",
    "# Load the text embeddings model\n",
    "from vertexai.preview.language_models import TextEmbeddingModel\n",
    "\n",
    "model = TextEmbeddingModel.from_pretrained(\"textembedding-gecko@001\")\n",
    "\n",
    "def get_embeddings_wrapper(texts):\n",
    "    embs = []\n",
    "    for i in tqdm.tqdm(range(0, len(texts), BATCH_SIZE)):\n",
    "        time.sleep(1)  # to avoid the quota error\n",
    "        result = model.get_embeddings(texts[i : i + BATCH_SIZE])\n",
    "        embs = embs + [e.values for e in result]\n",
    "    return embs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26c9fc2d-66ef-4f6b-a1e3-cfc14fda4b35",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df = pd.read_csv('df_exploded_2.csv')\n",
    "\n",
    "test_embeddings = get_embeddings_wrapper([\"Who is the best help to Harry Potter?\"])\n",
    "# Test query\n",
    "response = my_index_endpoint.find_neighbors(\n",
    "    deployed_index_id=DEPLOYED_INDEX_ID,\n",
    "    queries=test_embeddings,\n",
    "    num_neighbors=20,\n",
    ")\n",
    "\n",
    "# show the result\n",
    "import numpy as np\n",
    "\n",
    "for idx, neighbor in enumerate(response[0]):\n",
    "    id = np.int64(neighbor.id)\n",
    "    similar = df.query(\"id == @id\", engine=\"python\")\n",
    "    print(f\"{neighbor.distance:.4f} {similar.splitted_texts_chunks.values[0]}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2156c323-51a1-4170-975a-cc3be6f48ae6",
   "metadata": {},
   "source": [
    "### Get an existing Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fb220ea-3083-40e2-b7a3-d2c86ea8a7ba",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from google.cloud import aiplatform\n",
    "\n",
    "REGION = LOCATION = \"asia-southeast1\"\n",
    "\n",
    "aiplatform.init(project=PROJECT_ID, location=LOCATION)\n",
    "vertexai.init(project=PROJECT_ID, location=LOCATION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2551607-0f0d-479a-af00-5505f62fed1f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "my_index_name = my_index._gca_resource.name\n",
    "my_index_display_name = my_index.display_name\n",
    "my_index_id = my_index.name.split('/')[-1]\n",
    "\n",
    "my_index_endpoint_name = my_index_endpoint._gca_resource.name\n",
    "my_index_endpoint_display_name = my_index_endpoint.display_name\n",
    "my_index_endpoint_id = my_index_endpoint.name.split('/')[-1]\n",
    "my_index_endpoint_public_domain = my_index_endpoint.public_endpoint_domain_name\n",
    "\n",
    "my_index = aiplatform.MatchingEngineIndex(my_index_name)\n",
    "\n",
    "my_index_endpoint_id = my_index_endpoint_id\n",
    "\n",
    "# my_index_endpoint_id = \"[your-index-endpoint-id]\"  # @param {type:\"string\"}\n",
    "my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(my_index_endpoint_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cd7ded2-c8e9-4574-9756-5f0c68f85df4",
   "metadata": {},
   "source": [
    "### Querying the index created earlier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c35bc7b3-8366-40e0-ac2a-12eec665c1e7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain_google_vertexai import VertexAIEmbeddings\n",
    "\n",
    "embeddings = VertexAIEmbeddings(model=\"textembedding-gecko@001\", model_name=\"textembedding-gecko@001\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddc6af59-72da-4412-aed6-0e9089736ee4",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# this is embedding vector (should be created by calling the embeddings models)\n",
    "\n",
    "text = \"harry potter owl and the green colour boy\"\n",
    "\n",
    "test_embeddings = embeddings.embed_query(text)\n",
    "print(\"preview embeddings\",test_embeddings[0:2])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd987e48-ab41-444f-bee9-3004e57db53d",
   "metadata": {},
   "source": [
    "### Using Vector Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10c951d3-2d0c-4ba9-a687-9ef0d97318e7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# this setting is obtained from matching ending https://console.cloud.google.com/vertex-ai/locations/asia-southeast1/index-endpoints/3366088877738557440/deployed-indexes/vs_quickstart_deployed_02060053?project=jingle-project-414801\n",
    "\n",
    "from google.cloud import aiplatform_v1\n",
    "\n",
    "# Set variables for the current deployed index.\n",
    "API_ENDPOINT=my_index_endpoint_public_domain\n",
    "INDEX_ENDPOINT=my_index_endpoint_name\n",
    "\n",
    "indexendpoint_id=UNIQUE_PREFIX\n",
    "\n",
    "DEPLOYED_INDEX_ID=\"vs_feature_deployed_\" + indexendpoint_id\n",
    "neighbor_count = 3\n",
    "\n",
    "print(API_ENDPOINT)\n",
    "print(INDEX_ENDPOINT)\n",
    "print(DEPLOYED_INDEX_ID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fcb2e54-77fe-4a32-9835-8535e5917d17",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Configure Vector Search client\n",
    "client_options = {\n",
    "  \"api_endpoint\": API_ENDPOINT\n",
    "}\n",
    "vector_search_client = aiplatform_v1.MatchServiceClient(\n",
    "  client_options=client_options,\n",
    ")\n",
    "# Build FindNeighborsRequest object\n",
    "datapoint = aiplatform_v1.IndexDatapoint(\n",
    "  feature_vector=test_embeddings\n",
    ")\n",
    "\n",
    "query = aiplatform_v1.FindNeighborsRequest.Query(\n",
    "  datapoint=datapoint,\n",
    "  # The number of nearest neighbors to be retrieved\n",
    "  neighbor_count=neighbor_count\n",
    ")\n",
    "\n",
    "request = aiplatform_v1.FindNeighborsRequest(\n",
    "  index_endpoint=INDEX_ENDPOINT,\n",
    "  deployed_index_id=DEPLOYED_INDEX_ID,\n",
    "  # Request can have multiple queries\n",
    "  queries=[query],\n",
    "  return_full_datapoint=False,\n",
    ")\n",
    "\n",
    "# Execute the request\n",
    "response = vector_search_client.find_neighbors(request)\n",
    "\n",
    "df_new = pd.DataFrame()\n",
    "print('neighbor_count', neighbor_count)\n",
    "for i in range(0,neighbor_count):\n",
    "    x=response.nearest_neighbors[0]\n",
    "    \n",
    "    df_match = df.loc[df['id'] == int(x.neighbors[i].datapoint.datapoint_id) ]\n",
    "\n",
    "    # Append the matching rows to the new DataFrame\n",
    "    df_new = pd.concat([df_new, df_match])\n",
    "\n",
    "# Print the new DataFrame\n",
    "print(df_new)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05a56cc5-7562-44c4-8dc1-1dbe9ad5352d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_id_with_embedding_matching(test_embeddings) :\n",
    "    \n",
    "    datapoint = aiplatform_v1.IndexDatapoint(\n",
    "      feature_vector=test_embeddings\n",
    "    )\n",
    "    query = aiplatform_v1.FindNeighborsRequest.Query(\n",
    "      datapoint=datapoint,\n",
    "      # The number of nearest neighbors to be retrieved\n",
    "      neighbor_count=neighbor_count\n",
    "    )\n",
    "    request = aiplatform_v1.FindNeighborsRequest(\n",
    "      index_endpoint=INDEX_ENDPOINT,\n",
    "      deployed_index_id=DEPLOYED_INDEX_ID,\n",
    "      # Request can have multiple queries\n",
    "      queries=[query],\n",
    "      return_full_datapoint=False,\n",
    "    )\n",
    "\n",
    "    # Execute the request\n",
    "    response = vector_search_client.find_neighbors(request)\n",
    "    \n",
    "    df_new = pd.DataFrame()\n",
    "\n",
    "    for i in range(0,neighbor_count):\n",
    "        x=response.nearest_neighbors[0]\n",
    "\n",
    "        df_match = df.loc[df['id'] == int(x.neighbors[i].datapoint.datapoint_id) ]\n",
    "\n",
    "        # Append the matching rows to the new DataFrame\n",
    "        df_new = pd.concat([df_new, df_match])\n",
    "    \n",
    "    print(df_new)\n",
    "    i,j,k = df_new.index[0:3]\n",
    "    print(i,j,k)\n",
    "    \n",
    "    pagewise_texts_v1 = df_new.loc[i, 'pagewise_texts']\n",
    "    pagewise_texts_v2 = df_new.loc[j, 'pagewise_texts']\n",
    "    pagewise_texts_v3 = df_new.loc[k, 'pagewise_texts']\n",
    "    \n",
    "    splitted_texts_v1 = df_new.loc[i, 'splitted_texts']\n",
    "    splitted_texts_v2 = df_new.loc[j, 'pagewise_texts']\n",
    "    splitted_texts_v3 = df_new.loc[k, 'pagewise_texts']\n",
    "    \n",
    "    splitted_texts_chunks_v1 = df_new.loc[i, 'splitted_texts_chunks']\n",
    "    splitted_texts_chunks_v2 = df_new.loc[j, 'splitted_texts_chunks']\n",
    "    splitted_texts_chunks_v3 = df_new.loc[k, 'splitted_texts_chunks']\n",
    "    \n",
    "    page_id_v1 = df_new.loc[i, 'page_id'] \n",
    "    page_id_v2 = df_new.loc[j, 'page_id'] \n",
    "    page_id_v3 = df_new.loc[k, 'page_id'] \n",
    "    \n",
    "    return(pagewise_texts_v1,pagewise_texts_v2,pagewise_texts_v3,\n",
    "           splitted_texts_v1,splitted_texts_v2,splitted_texts_v3,\n",
    "           splitted_texts_chunks_v1,splitted_texts_chunks_v2,splitted_texts_chunks_v3,\n",
    "        page_id_v1,page_id_v2,page_id_v3,i,j,k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43dff9e2-1072-4deb-852e-50f3f857a6eb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "filename = f\"{FOLDER_NAME}/harry_potte_qa.csv\"\n",
    "df_qa = pd.read_csv(filename, sep =\"|\")\n",
    "\n",
    "df_qa.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7a45e77-0a55-4022-9f7a-1b13fb766ed7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_qa.columns"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe6a8963-e7cd-41af-8301-584eb5a399ff",
   "metadata": {},
   "source": [
    "### Export results into a CSV file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b4a83c7-f800-45aa-b5aa-ce29fbe66400",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "import csv\n",
    "\n",
    "with open('harry_potte_qa.csv', 'r') as input_file, open('harry_potte_qa_output.csv', 'w', newline='') as output_file:\n",
    "\n",
    "  # Create CSV reader and writer objects\n",
    "  reader = csv.reader(input_file, delimiter='|')\n",
    "  writer = csv.writer(output_file, delimiter='|')\n",
    "\n",
    "  # Read and write the header row\n",
    "  header = next(reader) + ['i','j','k','pagewise_texts_v1','pagewise_texts_v2','pagewise_texts_v3','splitted_texts_v1','splitted_texts_v2','splitted_texts_v3','splitted_texts_chunks_v1','splitted_texts_chunks_v2','splitted_texts_chunks_v3','page_id_v1','page_id_v2','page_id_v3']\n",
    "  writer.writerow(header)\n",
    "\n",
    "  # Loop through the remaining rows\n",
    "  for i, row in enumerate(reader):\n",
    "    question = row[0].split('|')[0]  # Use 'i' to access the correct element in the row\n",
    "    question_emb = embeddings.embed_query( question )\n",
    "    pagewise_texts_v1,pagewise_texts_v2,pagewise_texts_v3,splitted_texts_v1,splitted_texts_v2,splitted_texts_v3,splitted_texts_chunks_v1,splitted_texts_chunks_v2,splitted_texts_chunks_v3,page_id_v1,page_id_v2,page_id_v3,i,j,k = get_id_with_embedding_matching(question_emb) \n",
    "    \n",
    "    # print( i , question)\n",
    "    row_out = row + [i,j,k,pagewise_texts_v1,pagewise_texts_v2,pagewise_texts_v3,splitted_texts_v1,splitted_texts_v2,splitted_texts_v3,splitted_texts_chunks_v1,splitted_texts_chunks_v2,splitted_texts_chunks_v3,page_id_v1,page_id_v2,page_id_v3]\n",
    "    \n",
    "    # Write the row to the output file\n",
    "    writer.writerow(row_out)\n",
    "\n",
    "# Usage example:\n",
    "! head -n 2 harry_potte_qa_output.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a638e42d-4704-4c7d-8124-7889eac6b68b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "filename = f\"{FOLDER_NAME}/harry_potte_qa_output.csv\"\n",
    "df_qa = pd.read_csv(filename, sep =\"|\")\n",
    "\n",
    "df_qa.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b923bc10-ff7f-48c1-b144-23c25928c438",
   "metadata": {},
   "source": [
    "## Part 4: Trying Out Different AI Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba7ddb8c-43b0-4ac5-b935-8576d537cf7b",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "GCP_PROJECT=PROJECT_ID\n",
    "LOCATION = REGION = 'asia-southeast1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5537c3c-82eb-49d8-823f-c7bea7382221",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import vertexai\n",
    "from vertexai.preview.language_models import TextGenerationModel\n",
    "\n",
    "vertexai.init(project=PROJECT_ID, location=LOCATION)\n",
    "\n",
    "\n",
    "# @st.cache_resource\n",
    "def get_model():\n",
    "    generation_model = TextGenerationModel.from_pretrained(\"text-bison@002\")\n",
    "    return generation_model\n",
    "\n",
    "\n",
    "def get_text_generation(prompt=\"\", **parameters):\n",
    "    generation_model = get_model()\n",
    "    response = generation_model.predict(prompt=prompt, **parameters)\n",
    "\n",
    "    return response.text"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b06a05e6-64b2-4f2d-bcd0-dd60a56cd2ec",
   "metadata": {},
   "source": [
    "### Defining Functions for different models (Gemini, Unicorn, Bison32k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd6fdf26-46b3-459d-b831-27d72a865583",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import vertexai\n",
    "from vertexai.preview.generative_models import GenerativeModel, Part\n",
    "\n",
    "def generate(input_prompt):\n",
    "    model = GenerativeModel(\"gemini-ultra\")\n",
    "    responses = model.generate_content(\n",
    "        input_prompt ,\n",
    "    generation_config={\n",
    "        \"max_output_tokens\": 2048,\n",
    "        \"temperature\": 0.2,\n",
    "        \"top_p\": 1,\n",
    "        \"top_k\": 32\n",
    "    },\n",
    "        safety_settings=[],\n",
    "        stream=True,\n",
    "    )\n",
    "    \n",
    "    all_response  = []\n",
    "    \n",
    "    for response in responses:\n",
    "        # print(response.text, end=\"\")\n",
    "        all_response.append(response.text)\n",
    "    \n",
    "    # print (all_response)\n",
    "    \n",
    "    return(\" \".join(all_response))\n",
    "    \n",
    "\n",
    "def generate_pro(input_prompt):\n",
    "    model = GenerativeModel(\"gemini-pro\")\n",
    "    responses = model.generate_content(\n",
    "    input_prompt,\n",
    "    generation_config={\n",
    "        \"max_output_tokens\": 2048,\n",
    "        \"temperature\": 0.2,\n",
    "        \"top_p\": 1\n",
    "    },stream=True,)\n",
    "    \n",
    "    all_response  = []\n",
    "    \n",
    "    for response in responses:\n",
    "        all_response.append(response.text)\n",
    "    \n",
    "    # print (all_response)\n",
    "    \n",
    "    return(\" \".join(all_response))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9728825-9157-438b-9641-c298d7f7adb4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import vertexai\n",
    "from vertexai.language_models import TextGenerationModel\n",
    "\n",
    "vertexai.init(project=PROJECT_ID, location=LOCATION)\n",
    "parameters = {\n",
    "    \"candidate_count\": 1,\n",
    "    \"max_output_tokens\": 1024,\n",
    "    \"temperature\": 1,\n",
    "    \"top_k\": 40\n",
    "}\n",
    "\n",
    "def generate_palm_unicorn_v1(input_prompt):\n",
    "    \n",
    "    model = TextGenerationModel.from_pretrained(\"text-unicorn@001\")\n",
    "\n",
    "    response = model.predict(\n",
    "        input_prompt,\n",
    "        **parameters\n",
    "    )\n",
    "    print(f\"Response from Model: {response.text}\")\n",
    "    \n",
    "    return(response.text)\n",
    "\n",
    "def generate_palm_bison32k(input_prompt):\n",
    "    \n",
    "    model = TextGenerationModel.from_pretrained(\"text-bison-32k\")\n",
    "\n",
    "    response = model.predict(\n",
    "        input_prompt,\n",
    "        **parameters\n",
    "    )\n",
    "    print(f\"Response from Model: {response.text}\")\n",
    "    \n",
    "    return(response.text)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee960ef4-ebcb-4068-a997-0898d508faf6",
   "metadata": {},
   "source": [
    "### Read the Q&A file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fdc5924-f69e-4fa5-96b1-adf3b43721ed",
   "metadata": {},
   "source": [
    "#### This uses the file from Matching Engine which has questions and retrieved document results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dee6bf9f-c44b-4414-946e-b20f8d214c3b",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import re\n",
    "filename = f\"{FOLDER_NAME}/harry_potte_qa_output.csv\"\n",
    "df_qa = pd.read_csv(filename, sep =\"|\")\n",
    "\n",
    "System_Prompts = \"\"\" You are an expert in reading harry potter books, but only provide evidences from the information provided and do not use any other information\n",
    "so here are some search results : \n",
    "\"\"\"\n",
    "\n",
    "Question_Prompts = \"\"\" -- Based on information above help to answer following user question\n",
    "\"\"\"\n",
    "\n",
    "df_qa['combine_prompt_RAG1'] = System_Prompts + ' ' +df_qa['pagewise_texts_v1'] + ' Please answers the Question : '+ df_qa['Question'] \n",
    "df_qa['combine_prompt_RAG2'] = System_Prompts + ' ' +df_qa['pagewise_texts_v2'] + ' Please answers the Question : '+ df_qa['Question'] \n",
    "df_qa['combine_prompt_RAG3'] = System_Prompts + ' ' +df_qa['pagewise_texts_v3'] + ' Please answers the Question : '+ df_qa['Question'] \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dafb13b-7d55-4803-97de-d783045351b7",
   "metadata": {},
   "source": [
    "### Submit all the questions into different AI Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dfc0578-fddd-4378-9da5-bdb0616f81de",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "for i in range(0, len(df_qa)):\n",
    "\n",
    "\n",
    "    clean_text1 = re.sub(r'[^\\w\\s;]', '', df_qa.loc[i,'combine_prompt_RAG1'])\n",
    "    clean_text2 = re.sub(r'[^\\w\\s;]', '', df_qa.loc[i,'combine_prompt_RAG2'])\n",
    "    clean_text3 = re.sub(r'[^\\w\\s;]', '', df_qa.loc[i,'combine_prompt_RAG3'])\n",
    "\n",
    "    if i<=1000:\n",
    "        # df['Gemini_ultra_model_output'][i] = generate(df['combine_prompt'][i])\n",
    "        print(\"iteration #\", i, \"test\")\n",
    "        if i==32 : \n",
    "            print(\"iteration #\", i, \"test\", clean_text1, clean_text2, clean_text3)\n",
    "    \n",
    "    try:\n",
    "        df_qa.loc[i, \"Gemini_pro_model_output_v1\"] = generate_pro(clean_text1)\n",
    "        df_qa.loc[i, \"Gemini_pro_model_output_v2\"] = generate_pro(clean_text2)\n",
    "        df_qa.loc[i, \"Gemini_pro_model_output_v3\"] = generate_pro(clean_text3)\n",
    "    except :\n",
    "        print(\"Prompt error at gemini \", i)\n",
    "        df_qa.loc[i, \"Gemini_pro_model_output_v1\"] = \"Prompt failed \"\n",
    "        df_qa.loc[i, \"Gemini_pro_model_output_v2\"] = \"Prompt failed \"\n",
    "        df_qa.loc[i, \"Gemini_pro_model_output_v3\"] = \"Prompt failed \"\n",
    "\n",
    "    try:\n",
    "        df_qa.loc[i, \"palm_bison32k_output_v1\"] = generate_palm_bison32k(df_qa.loc[i,'combine_prompt_RAG1'])\n",
    "        df_qa.loc[i, \"palm_bison32k_output_v2\"] = generate_palm_bison32k(df_qa.loc[i,'combine_prompt_RAG2'])\n",
    "        df_qa.loc[i, \"palm_bison32k_output_v3\"] = generate_palm_bison32k(df_qa.loc[i,'combine_prompt_RAG3'])\n",
    "    except ZeroDivisionError:\n",
    "        print(\"Prompt error at palm \", i)\n",
    "        df_qa.loc[i, \"palm_bison32k_output_v1\"] = \"Prompt failed \"\n",
    "        df_qa.loc[i, \"palm_bison32k_output_v2\"] = \"Prompt failed \"\n",
    "        df_qa.loc[i, \"palm_bison32k_output_v3\"] = \"Prompt failed \"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a245d4dc-1342-4393-b9be-511f7ee971c1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_qa.columns"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea0106fd-f67b-4c89-ac95-c5caa5fee2f9",
   "metadata": {},
   "source": [
    "### Output all Questions & Answers from all AI Models into a CSV file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7c44c34-ed52-443d-bdf0-d4d19662c64c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Delete the 'col2' column\n",
    "df_qa = df_qa.drop('combine_prompt_RAG1', axis=1)\n",
    "df_qa = df_qa.drop('combine_prompt_RAG2', axis=1)\n",
    "df_qa = df_qa.drop('combine_prompt_RAG3', axis=1)\n",
    "\n",
    "output1 = f\"{FOLDER_NAME}/results/harry_potte_qa_model_out.csv\"\n",
    "\n",
    "df_qa.to_csv(output1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cf83e7c-ec48-4b58-8d8c-464bf902afb7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_qa.head(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34e861c6-76a3-4cc3-bfb0-767876c9feac",
   "metadata": {},
   "source": [
    "### The response is not ideal due to the poor search setup. This can be finetuned further for better results."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "147b83ea-6d13-4099-995a-7416ea71d9da",
   "metadata": {
    "tags": []
   },
   "source": [
    "# End of Lab"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "conda-base-py",
   "name": "workbench-notebooks.m121",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m121"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel) (Local)",
   "language": "python",
   "name": "conda-base-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}