rag_pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import getpass\n",
    "os.environ['LANGCHAIN_TRACING_V2'] = 'true'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdf_path = 'pdfs'\n",
    "pdf_files = [os.path.join(pdf_path, f) for f in os.listdir(pdf_path) if f.endswith('.pdf')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdf_files = [ # always dict with keys path, title, authors\n",
    "     {'path': 'pdfs/1706.03762v7.pdf', 'title': 'Attention Is All You Need', 'authors': 'Vaswani, Ashish and Shazeer, Noam and Parmar, Niki and Uszkoreit, Jakob and Jones, Llion and Gomez, Aidan N and Kaiser, Lukasz and Polosukhin, Illia'},\n",
    "     {'path': 'pdfs/srivastava14a.pdf', 'title': 'Dropout: A Simple Way to Prevent Neural Networks from Overfitting', 'authors': 'Srivastava, Nitish and Hinton, Geoffrey and Krizhevsky, Alex and Sutskever, Ilya and Salakhutdinov, Ruslan'}\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.document_loaders import PyPDFLoader\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "from langchain_chroma import Chroma\n",
    "\n",
    "gather_pages = []\n",
    "for pdf_file in pdf_files:\n",
    "    # PdPDFLoader excepts strings as file paths, allows to extract images via extract_images = True\n",
    "    loader = PyPDFLoader(pdf_file[\"path\"])\n",
    "    pages = loader.load_and_split()\n",
    "    metadata = f\"Title: {pdf_file['title']}\\nAuthors: {pdf_file['authors']}\"\n",
    "    for page in pages:\n",
    "        page.metadata[\"source\"] = metadata\n",
    "    gather_pages.extend(pages)\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)\n",
    "splits = text_splitter.split_documents(gather_pages)\n",
    "\n",
    "vectorstore = Chroma.from_documents(documents=splits, embedding=OpenAIEmbeddings(model=\"text-embedding-3-small\"), persist_directory=\"db/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "vectorstore = Chroma(embedding_function=OpenAIEmbeddings(model=\"text-embedding-3-small\"), persist_directory=\"db/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_context(context):\n",
    "    \"\"\"Format the context for the prompt.\n",
    "    Args:\n",
    "        context: list of tuples (document, score) from the retriever\n",
    "    Returns:\n",
    "        str: formatted context\n",
    "    \"\"\"\n",
    "    output = \"\"\n",
    "    for doc in context:\n",
    "        output += doc[0].metadata[\"source\"] + f\"\\nPage: {doc[0].metadata['page']}\" +\\\n",
    "        f\"\\nPage content: {doc[0].page_content}\\n\\n\"\n",
    "    return output\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reproduction is an important concept in computer science, particularly in the context of evolutionary algorithms and neural network training. The principles of sexual reproduction, such as gene mixing and mutation, are often applied to optimize algorithms, enhancing robustness and adaptability in machine learning models. This concept is exemplified in the work that discusses dropout techniques in neural networks, drawing parallels between biological evolution and algorithmic improvement to prevent overfitting and enhance performance^1.\n",
      "\n",
      "---\n",
      "^1 Srivastava, Nitish et al. \"Dropout: A Simple Way to Prevent Neural Networks from Overfitting,\" page 3.\n"
     ]
    }
   ],
   "source": [
    "query = \"Is reproduction an important concept in computer science?\"\n",
    "num_context_chunks = 10\n",
    "results = vectorstore.similarity_search_with_score(\n",
    "    query , k=num_context_chunks\n",
    ")\n",
    "\n",
    "formatted_context = format_context(results)\n",
    "prompt = \"Systemprompt:\\nYou are an assistant for question-answering tasks. Use the following pieces of \" + \\\n",
    "\"retrieved context to answer the question and give reference to the publication title, authors \" +\\\n",
    "\"and page as footnotes but without duplicate references. If you don't know the answer, just say \"+ \\\n",
    "\"that you don't know. Please keep the answer concise and precise.\\n\\n\" +\\\n",
    "f\"Question:\\n{query}\\n\\n\" +\\\n",
    "f\"Context:\\n{formatted_context}\"\n",
    "\n",
    "answer = llm.call_as_llm(prompt)\n",
    "print(answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "File pdfs/gal16.pdf uploaded successfully to 664418962820-team-air-1/gal16.pdf with metadata\n"
     ]
    }
   ],
   "source": [
    "import boto3\n",
    "import os\n",
    "\n",
    "\n",
    "s3 = boto3.client('s3')\n",
    "\n",
    "def upload_file_to_s3(file_path, bucket_name, object_name=None, metadata=None):\n",
    "    if object_name is None:\n",
    "        object_name = file_path.split('/')[-1]\n",
    "    \n",
    "    # Ensure metadata is a dictionary\n",
    "    if metadata is None:\n",
    "        metadata = {}\n",
    "    \n",
    "    # Prepare the ExtraArgs parameter\n",
    "    extra_args = {\n",
    "        'Metadata': metadata,\n",
    "        'ContentType': 'application/pdf'  # Adjust this based on your file type\n",
    "    }\n",
    "    \n",
    "    try:\n",
    "        s3.upload_file(file_path, bucket_name, object_name, ExtraArgs=extra_args)\n",
    "        print(f\"File {file_path} uploaded successfully to {bucket_name}/{object_name} with metadata\")\n",
    "    except Exception as e:\n",
    "        print(f\"Error uploading file: {e}\")\n",
    "\n",
    "metadata = {\n",
    "    'title': 'Dropout as a bayesian approximation: Representing model uncertainty in deep learning',\n",
    "    'authors': 'Gal, Yarin and Ghahramani, Zoubin',\n",
    "    'year': '2016',\n",
    "}\n",
    "\n",
    "# Example usage\n",
    "upload_file_to_s3(\n",
    "    file_path='pdfs/gal16.pdf', bucket_name='664418962820-team-air-1', metadata=metadata\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ingestion job started. Job ID: 5NYBRBAHWU\n",
      "Re-sync in progress. Current status: STARTING\n",
      "Re-sync completed successfully.\n"
     ]
    }
   ],
   "source": [
    "import boto3\n",
    "import time\n",
    "\n",
    "def resync_bedrock_knowledge_base(knowledge_base_id, data_source_id):\n",
    "    bedrock = boto3.client('bedrock-agent', region_name='eu-central-1')\n",
    "    try:\n",
    "        # Start a new ingestion job\n",
    "        response = bedrock.start_ingestion_job(\n",
    "            knowledgeBaseId=knowledge_base_id,\n",
    "            dataSourceId=data_source_id,\n",
    "            description='Re-sync knowledge base with S3 bucket'\n",
    "        )\n",
    "        \n",
    "        job_id = response['ingestionJob']['ingestionJobId']\n",
    "        print(f\"Ingestion job started. Job ID: {job_id}\")\n",
    "        \n",
    "        # Wait for the job to complete\n",
    "        while True:\n",
    "            job_status = bedrock.get_ingestion_job(\n",
    "                knowledgeBaseId=knowledge_base_id,\n",
    "                dataSourceId=data_source_id,\n",
    "                ingestionJobId=job_id\n",
    "            )['ingestionJob']['status']\n",
    "            \n",
    "            if job_status == 'COMPLETE':\n",
    "                print(\"Re-sync completed successfully.\")\n",
    "                break\n",
    "            elif job_status in ['FAILED', 'CANCELLED']:\n",
    "                print(f\"Re-sync failed. Status: {job_status}\")\n",
    "                break\n",
    "            else:\n",
    "                print(f\"Re-sync in progress. Current status: {job_status}\")\n",
    "                time.sleep(3)  # Wait for 30 seconds before checking again\n",
    "                \n",
    "    except Exception as e:\n",
    "        print(f\"An error occurred: {str(e)}\")\n",
    "\n",
    "# Example usage\n",
    "resync_bedrock_knowledge_base(knowledge_base_id='4FYUGYITNF', data_source_id='PIKXJTPZTB')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain",
   "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.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}