From e3432ce680fa4dc9394f2f3c648c2d950643fb40 Mon Sep 17 00:00:00 2001 From: jonfritz <134336691+jonfritz@users.noreply.github.com> Date: Fri, 17 Jan 2025 11:15:40 -0800 Subject: [PATCH] Add notebook for Aryn blog post (#380) * Create test * Add files via upload * Rename aryn-elasticsearch-blog-demo-clean copy.ipynb to aryn-elasticsearch-blog-dataprep.ipynb * Create README.md * Delete notebooks/integrations/aryn/test * Add files via upload Update with pip install * Update aryn-elasticsearch-blog-dataprep.ipynb Update from feedback * Update aryn-elasticsearch-blog-dataprep.ipynb Add placeholder API key * Update aryn-elasticsearch-blog-dataprep.ipynb Add comma * Update aryn-elasticsearch-blog-dataprep.ipynb Formatting * Update aryn-elasticsearch-blog-dataprep.ipynb Formatting * Update aryn-elasticsearch-blog-dataprep.ipynb Add API key placeholder * Update aryn-elasticsearch-blog-dataprep.ipynb Formatting * fix formatting Signed-off-by: Henry Lindeman * Create aryn-elasticsearch-RAG-data-preparation-demo * Rename supporting-blog-content/aryn-elasticsearch-RAG-data-preparation-demo to supporting-blog-content/Aryn-elasticsearch-RAG-data-preparation-demo/test.md * Add files via upload * Delete supporting-blog-content/Aryn-elasticsearch-RAG-data-preparation-demo/test.md * Delete notebooks/integrations/aryn/aryn-elasticsearch-blog-dataprep.ipynb * Update README.md Add link --------- Signed-off-by: Henry Lindeman Co-authored-by: Henry Lindeman --- notebooks/integrations/aryn/README.md | 3 + .../aryn-elasticsearch-blog-dataprep.ipynb | 208 ++++++++++++++++++ 2 files changed, 211 insertions(+) create mode 100644 notebooks/integrations/aryn/README.md create mode 100644 supporting-blog-content/Aryn-elasticsearch-RAG-data-preparation-demo/aryn-elasticsearch-blog-dataprep.ipynb diff --git a/notebooks/integrations/aryn/README.md b/notebooks/integrations/aryn/README.md new file mode 100644 index 00000000..bfe8c35e --- /dev/null +++ b/notebooks/integrations/aryn/README.md @@ -0,0 +1,3 @@ +This folder contains examples showing how to prepare data using Aryn Sycamore and load into Elasticsearch for RAG and GenAI use cases. + +The notebook for the Aryn Elasticsearch blog example is [here](https://github.com/elastic/elasticsearch-labs/blob/main/supporting-blog-content/Aryn-elasticsearch-RAG-data-preparation-demo/aryn-elasticsearch-blog-dataprep.ipynb). diff --git a/supporting-blog-content/Aryn-elasticsearch-RAG-data-preparation-demo/aryn-elasticsearch-blog-dataprep.ipynb b/supporting-blog-content/Aryn-elasticsearch-RAG-data-preparation-demo/aryn-elasticsearch-blog-dataprep.ipynb new file mode 100644 index 00000000..4a3c0bc0 --- /dev/null +++ b/supporting-blog-content/Aryn-elasticsearch-RAG-data-preparation-demo/aryn-elasticsearch-blog-dataprep.ipynb @@ -0,0 +1,208 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "a8f66d95-a9c4-40f1-8cf8-19795653c3f3", + "metadata": {}, + "outputs": [], + "source": [ + "!pip install sycamore-ai[elasticsearch]\n", + "# Install the Sycamore document ETL library: https://github.com/aryn-ai/sycamore" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "60b49e1c-7055-4534-ac09-8b7ab45086d4", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import sycamore\n", + "from sycamore.context import ExecMode\n", + "from sycamore.transforms.partition import ArynPartitioner\n", + "from sycamore.transforms.extract_schema import LLMPropertyExtractor\n", + "from sycamore.transforms.summarize_images import SummarizeImages, LLMImageSummarizer\n", + "from sycamore.transforms.standardizer import (\n", + " USStateStandardizer,\n", + " DateTimeStandardizer,\n", + " ignore_errors,\n", + ")\n", + "from sycamore.transforms.merge_elements import GreedySectionMerger\n", + "from sycamore.functions.tokenizer import HuggingFaceTokenizer\n", + "from sycamore.transforms.embed import SentenceTransformerEmbedder\n", + "from sycamore.llms import OpenAI, OpenAIModels\n", + "\n", + "import pyarrow.fs\n", + "\n", + "llm = OpenAI(OpenAIModels.GPT_4O_MINI)\n", + "os.environ[\"ARYN_API_KEY\"] = \"\"\n", + "\n", + "paths = [\"s3://aryn-public/ntsb/\"]\n", + "\n", + "context = sycamore.init()\n", + "# Add exec_mode=ExecMode.LOCAL to .init to run without Ray\n", + "docset = context.read.binary(paths=paths, binary_format=\"pdf\")\n", + "docset = docset.materialize(\n", + " path=\"./elasticsearch-tutorial/downloaded-docset\",\n", + " source_mode=sycamore.MATERIALIZE_USE_STORED,\n", + ")\n", + "# Make sure your Aryn token is accessible in the environment variable ARYN_API_KEY\n", + "partitioned_docset = docset.partition(\n", + " partitioner=ArynPartitioner(extract_table_structure=True, extract_images=True)\n", + ").materialize(\n", + " path=\"./elasticsearch-tutorial/partitioned-docset\",\n", + " source_mode=sycamore.MATERIALIZE_USE_STORED,\n", + ")\n", + "partitioned_docset.execute()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a755a09e-1622-400b-8b75-b3bad2981b5f", + "metadata": {}, + "outputs": [], + "source": [ + "schema = {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"accidentNumber\": {\"type\": \"string\"},\n", + " \"dateAndTime\": {\"type\": \"date\"},\n", + " \"location\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"US State where the incident occured\",\n", + " },\n", + " \"aircraft\": {\"type\": \"string\"},\n", + " \"aircraftDamage\": {\"type\": \"string\"},\n", + " \"injuries\": {\"type\": \"string\"},\n", + " \"definingEvent\": {\"type\": \"string\"},\n", + " },\n", + " \"required\": [\"accidentNumber\", \"dateAndTime\", \"location\", \"aircraft\"],\n", + "}\n", + "\n", + "schema_name = \"FlightAccidentReport\"\n", + "property_extractor = LLMPropertyExtractor(\n", + " llm=llm, num_of_elements=20, schema_name=schema_name, schema=schema\n", + ")\n", + "\n", + "enriched_docset = (\n", + " partitioned_docset\n", + " # Extracts the properties based on the schema defined\n", + " .extract_properties(property_extractor=property_extractor)\n", + " # Summarizes images that were extracted using an LLM\n", + " .transform(SummarizeImages, summarizer=LLMImageSummarizer(llm=llm))\n", + ")\n", + "\n", + "formatted_docset = (\n", + " enriched_docset\n", + " # Converts state abbreviations to their full names.\n", + " .map(\n", + " lambda doc: ignore_errors(\n", + " doc, USStateStandardizer, [\"properties\", \"entity\", \"location\"]\n", + " )\n", + " )\n", + " # Converts datetime into a common format\n", + " .map(\n", + " lambda doc: ignore_errors(\n", + " doc, DateTimeStandardizer, [\"properties\", \"entity\", \"dateAndTime\"]\n", + " )\n", + " )\n", + ")\n", + "\n", + "\n", + "merger = GreedySectionMerger(\n", + " tokenizer=HuggingFaceTokenizer(\"sentence-transformers/all-MiniLM-L6-v2\"),\n", + " max_tokens=512,\n", + ")\n", + "chunked_docset = formatted_docset.merge(merger=merger)\n", + "\n", + "model_name = \"thenlper/gte-small\"\n", + "\n", + "embedded_docset = (\n", + " chunked_docset.spread_properties([\"entity\", \"path\"])\n", + " .explode()\n", + " .embed(\n", + " embedder=SentenceTransformerEmbedder(batch_size=10_000, model_name=model_name)\n", + " )\n", + ")\n", + "\n", + "embedded_docset = embedded_docset.materialize(\n", + " path=\"./elasticsearch-tutorial/embedded-docset\",\n", + " source_mode=sycamore.MATERIALIZE_USE_STORED,\n", + ")\n", + "embedded_docset.execute()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b9321d7e-e812-41ac-8030-3db80c2147ec", + "metadata": {}, + "outputs": [], + "source": [ + "# Write to a persistent Elasticsearch Index. Note: You must have a specified elasticsearch instance running for this to work.\n", + "# For more information on how to set one up, refer to https://www.elastic.co/guide/en/elasticsearch/reference/current/install-elasticsearch.html\n", + "\n", + "url = \"http://localhost:9200\"\n", + "index_name = \"aryn-demo\"\n", + "embedded_ds.write.elasticsearch(\n", + " url=url,\n", + " index_name=index_name,\n", + " es_client_args={\"basic_auth\": (\"\", os.getenv(\"ELASTIC_PASSWORD\"))},\n", + " mappings={\n", + " \"properties\": {\n", + " \"embeddings\": {\n", + " \"type\": \"dense_vector\",\n", + " \"dims\": dimensions,\n", + " \"index\": True,\n", + " \"similarity\": \"cosine\",\n", + " },\n", + " \"properties\": {\"type\": \"object\"},\n", + " }\n", + " },\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "52970be4-7bac-455b-bcd0-868130ac61fd", + "metadata": {}, + "outputs": [], + "source": [ + "# Verify data has been loaded using DocSet Query to retrieve chunks\n", + "query_params = {\"match_all\": {}}\n", + "query_docs = ctx.read.elasticsearch(\n", + " url=url,\n", + " index_name=index_name,\n", + " query=query_params,\n", + " es_client_args={\"basic_auth\": (\"\", os.getenv(\"ELASTIC_PASSWORD\"))},\n", + ")\n", + "query_docs.show(show_embedding=False)" + ] + } + ], + "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.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}