sycamore_local_dev_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1666a1c-8c22-4643-8e15-a9a9f863893f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import subprocess\n",
    "import sys\n",
    "\n",
    "from pathlib import Path\n",
    "from IPython.display import display, Image\n",
    "from IPython.display import IFrame\n",
    "from PIL import Image as PImage\n",
    "\n",
    "import sycamore\n",
    "from sycamore.data import Document\n",
    "from sycamore.transforms.embed import SentenceTransformerEmbedder\n",
    "from sycamore.transforms.extract_entity import OpenAIEntityExtractor\n",
    "from sycamore.llms import OpenAIModels, OpenAI, LLM\n",
    "from sycamore.transforms.partition import UnstructuredPdfPartitioner, HtmlPartitioner\n",
    "from sycamore.llms.prompts.default_prompts import TEXT_SUMMARIZER_GUIDANCE_PROMPT_CHAT\n",
    "from sycamore.transforms.summarize import Summarizer\n",
    "from sycamore.transforms.extract_table import TextractTableExtractor\n",
    "from sycamore.functions.document import split_and_convert_to_image, DrawBoxes\n",
    "from sycamore.tests.config import TEST_DIR\n",
    "from sycamore.transforms.merge_elements import GreedyTextElementMerger\n",
    "from sycamore.functions.tokenizer import HuggingFaceTokenizer\n",
    "from sycamore.scans.file_scan import JsonManifestMetadataProvider"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6543822-4e34-4188-8887-a50e2251b2d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "work_dir = \"tmp/sycamore/data\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9632062e-b22e-4a23-b226-798961a05d29",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(work_dir, exist_ok = True)\n",
    "metadata = {}\n",
    "for f in [\"1706.03762.pdf\", \"2306.07303.pdf\"]:\n",
    "    path = os.path.join(work_dir, f)\n",
    "    url = os.path.join(\"https://arxiv.org/pdf\", f)\n",
    "    if not Path(path).is_file():\n",
    "        print(\"Downloading {} to {}\".format(url, path))\n",
    "        subprocess.run([\"curl\", \"-o\", path, url])\n",
    "    metadata[path] = { \"_location\": url }\n",
    "\n",
    "manifest_path = os.path.join(work_dir, \"manifest.json\")\n",
    "with open(manifest_path, \"w\") as f:\n",
    "    json.dump(metadata, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08cff60d-167a-465d-a135-0ee525fca1f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "IFrame(str(metadata[os.path.join(work_dir, \"1706.03762.pdf\")][\"_location\"]), width=700, height=600)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94e33e3b-2d95-408c-a09f-0786beede9ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "IFrame(str(metadata[os.path.join(work_dir, \"2306.07303.pdf\")][\"_location\"]), width=700, height=600)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "751e0849-49b0-4a54-8615-cf74e68217c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "openai_llm = OpenAI(OpenAIModels.GPT_3_5_TURBO.value)\n",
    "\n",
    "context = sycamore.init()\n",
    "pdf_docset = context.read.binary(work_dir, binary_format=\"pdf\", metadata_provider=JsonManifestMetadataProvider(manifest_path))\n",
    "\n",
    "pdf_docset.show(show_binary = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6ea4183-4b58-4278-83db-c16420b1cd9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: these fonts aren't correct, but are close enough for the visualization\n",
    "if os.path.isfile(\"/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf\"):\n",
    "    font_path = \"LiberationSans-Regular\"\n",
    "else:\n",
    "    print(\"Using default Arial font, which should work on MacOS and Windows\")\n",
    "    font_path = \"Arial.ttf\"\n",
    "\n",
    "def filter_func(doc: Document) -> bool:\n",
    "    return doc.properties[\"page_number\"] == 1\n",
    "\n",
    "partitioned_docset = pdf_docset.partition(partitioner=UnstructuredPdfPartitioner())\n",
    "visualized_docset = (partitioned_docset\n",
    "              .flat_map(split_and_convert_to_image)\n",
    "              .map_batch(DrawBoxes, f_constructor_args=[font_path])\n",
    "              .filter(filter_func))\n",
    "\n",
    "for doc in visualized_docset.take(2):\n",
    "    display(Image(doc.binary_representation, height=500, width=500))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3af55ed-fb1a-47f2-afac-6f10df824f60",
   "metadata": {},
   "outputs": [],
   "source": [
    " title_context_template = \"\"\"\n",
    "    ELEMENT 1: Jupiter's Moons\n",
    "    ELEMENT 2: Ganymede 2020\n",
    "    ELEMENT 3: by Audi Lauper and Serena K. Goldberg. 2011\n",
    "    ELEMENT 4: From Wikipedia, the free encyclopedia\n",
    "    ELEMENT 5: Ganymede, or Jupiter III, is the largest and most massive natural satellite of Jupiter as well as in the Solar System, being a planetary-mass moon. It is the largest Solar System object without an atmosphere, despite being the only moon of the Solar System with a magnetic field. Like Titan, it is larger than the planet Mercury, but has somewhat less surface gravity than Mercury, Io or the Moon.\n",
    "    =========\n",
    "    \"Ganymede 2020\"\n",
    "\n",
    "    ELEMENT 1: FLAVR: Flow-Agnostic Video Representations for Fast Frame Interpolation\n",
    "    ELEMENT 2: Tarun Kalluri * UCSD\n",
    "    ELEMENT 3: Deepak Pathak CMU\n",
    "    ELEMENT 4: Manmohan Chandraker UCSD\n",
    "    ELEMENT 5: Du Tran Facebook AI\n",
    "    ELEMENT 6: https://tarun005.github.io/FLAVR/\n",
    "    ELEMENT 7: 2 2 0 2\n",
    "    ELEMENT 8: b e F 4 2\n",
    "    ELEMENT 9: ]\n",
    "    ELEMENT 10: V C . s c [\n",
    "    ========\n",
    "    \"FLAVR: Flow-Agnostic Video Representations for Fast Frame Interpolation\"\n",
    "    \n",
    "    \"\"\"\n",
    "author_context_template = \"\"\"\n",
    "    ELEMENT 1: Jupiter's Moons\n",
    "    ELEMENT 2: Ganymede 2020\n",
    "    ELEMENT 3: by Audi Lauper and Serena K. Goldberg. 2011\n",
    "    ELEMENT 4: From Wikipedia, the free encyclopedia\n",
    "    ELEMENT 5: Ganymede, or Jupiter III, is the largest and most massive natural satellite of Jupiter as well as in the Solar System, being a planetary-mass moon. It is the largest Solar System object without an atmosphere, despite being the only moon of the Solar System with a magnetic field. Like Titan, it is larger than the planet Mercury, but has somewhat less surface gravity than Mercury, Io or the Moon.\n",
    "    =========\n",
    "    Audi Laupe, Serena K. Goldberg\n",
    "\n",
    "    ELEMENT 1: FLAVR: Flow-Agnostic Video Representations for Fast Frame Interpolation\n",
    "    ELEMENT 2: Tarun Kalluri * UCSD\n",
    "    ELEMENT 3: Deepak Pathak CMU\n",
    "    ELEMENT 4: Manmohan Chandraker UCSD\n",
    "    ELEMENT 5: Du Tran Facebook AI\n",
    "    ELEMENT 6: https://tarun005.github.io/FLAVR/\n",
    "    ELEMENT 7: 2 2 0 2\n",
    "    ELEMENT 8: b e F 4 2\n",
    "    ELEMENT 9: ]\n",
    "    ELEMENT 10: V C . s c [\n",
    "    ========\n",
    "    Tarun Kalluri, Deepak Pathak, Manmohan Chandraker, Du Tran\n",
    "\n",
    "  \"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7389a4ef-fa9d-4009-8d1e-0ba3ee99ca4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "entity_docset = (partitioned_docset\n",
    "                 .extract_entity(entity_extractor=OpenAIEntityExtractor(\"title\", llm=openai_llm, prompt_template=title_context_template))\n",
    "                 .extract_entity(entity_extractor=OpenAIEntityExtractor(\"authors\", llm=openai_llm, prompt_template=author_context_template)))\n",
    "\n",
    "entity_docset = entity_docset.spread_properties([\"title\", \"authors\"])\n",
    "\n",
    "entity_docset.show(show_binary = False, show_elements=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8bd0d95-519a-40a7-b1ab-ce404b42190e",
   "metadata": {},
   "outputs": [],
   "source": [
    "merged_docset = entity_docset.merge(GreedyTextElementMerger(tokenizer=HuggingFaceTokenizer(\"sentence-transformers/all-MiniLM-L6-v2\"), max_tokens=512))\n",
    "merged_docset.show(show_binary = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e22de01-756c-479d-bbd4-10232ef6caec",
   "metadata": {},
   "outputs": [],
   "source": [
    "exploded_docset = merged_docset.explode()\n",
    "exploded_docset.show(show_binary = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b37e4a6-ffd5-4bb7-931b-db1848430266",
   "metadata": {},
   "outputs": [],
   "source": [
    "st_embed_docset = (exploded_docset\n",
    "              .embed(embedder=SentenceTransformerEmbedder(batch_size=100, model_name=\"sentence-transformers/all-MiniLM-L6-v2\")))\n",
    "st_embed_docset.show(show_binary = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2abfa017-5661-4fe7-b04c-ca130bbfc9f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "index = \"local_development_example_index_withentity\" # You can change this to something else if you'd like\n",
    "\n",
    "os_client_args = {\n",
    "        \"hosts\": [{\"host\": \"localhost\", \"port\": 9200}],\n",
    "        \"http_compress\": True,\n",
    "        \"http_auth\": (\"admin\", \"admin\"),\n",
    "        \"use_ssl\": False,\n",
    "        \"verify_certs\": False,\n",
    "        \"ssl_assert_hostname\": False,\n",
    "        \"ssl_show_warn\": False,\n",
    "        \"timeout\": 120,\n",
    "    }\n",
    "\n",
    "index_settings =  {\n",
    "        \"body\": {\n",
    "            \"settings\": {\"index.knn\": True, \"number_of_shards\": 5, \"number_of_replicas\": 1},\n",
    "            \"mappings\": {\n",
    "                \"properties\": {\n",
    "                    \"embedding\": {\n",
    "                        \"dimension\": 384,\n",
    "                        \"method\": {\"engine\": \"nmslib\", \"space_type\": \"l2\", \"name\": \"hnsw\", \"parameters\": {}},\n",
    "                        \"type\": \"knn_vector\",\n",
    "                    }\n",
    "                }\n",
    "            },\n",
    "        }\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e8004aa-583c-4574-9bd7-fe092ffa737e",
   "metadata": {},
   "outputs": [],
   "source": [
    "st_embed_docset.write.opensearch(os_client_args=os_client_args, index_name=index, index_settings=index_settings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f08e3d7-183f-4014-9d31-7753e551954e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Visit http://localhost:3000 and use the\", index, \" index to query these results in the UI\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}