vmware · yonitoo
diff --git a/examples/embed-ingest-job-example/20_clean_and_embed_json_data.py b/examples/embed-ingest-job-example/20_clean_and_embed_json_data.py
@@ -0,0 +1,104 @@
+# Copyright 2021-2024 VMware, Inc.
+# SPDX-License-Identifier: Apache-2.0
+import json
+import logging
+import pathlib
+import re
+
+import nltk
+from config import DOCUMENTS_JSON_FILE_LOCATION
+from config import EMBEDDINGS_PKL_FILE_LOCATION
+from nltk.corpus import stopwords
+from nltk.stem import WordNetLemmatizer
+from sentence_transformers import SentenceTransformer
+from vdk.api.job_input import IJobInput
+
+log = logging.getLogger(__name__)
+
+
+def clean_text(text):
+    """
+    Prepares text for NLP tasks (embedding and RAG) by standardizing its form. It focuses on retaining
+    meaningful words and achieving consistency in their representation. This involves
+    converting to lowercase (uniformity), removing punctuation and stopwords
+    (focusing on relevant words), and lemmatization (reducing words to their base form).
+    Such preprocessing is crucial for effective NLP analysis.
+
+    :param text: A string containing the text to be processed.
+    :return: The processed text as a string.
+    """
+    text = text.lower()
+    # remove punctuation and special characters
+    text = re.sub(r"[^\w\s]", "", text)
+    # remove stopwords and lemmatize
+    stop_words = set(stopwords.words("english"))
+    lemmatizer = WordNetLemmatizer()
+    text = " ".join(
+        [lemmatizer.lemmatize(word) for word in text.split() if word not in stop_words]
+    )
+    return text
+
+
+def load_and_clean_documents(json_file_path):
+    cleaned_documents = []
+    with open(json_file_path, encoding="utf-8") as file:
+        documents = json.load(file)
+
+    for doc in documents:
+        if "data" in doc:
+            cleaned_text = clean_text(doc["data"])
+            cleaned_documents.append([cleaned_text])
+
+    print(len(cleaned_documents))
+    return cleaned_documents
+
+
+def embed_documents_in_batches(documents):
+    # the model card: https://huggingface.co/sentence-transformers/all-mpnet-base-v2
+    model = SentenceTransformer("all-mpnet-base-v2")
+    total = len(documents)
+    log.info(f"total: {total}")
+    embeddings = []
+    for start_index in range(0, total):
+        # the resources are not enough to batch 2 documents at a time, so the batch = 1 doc
+        batch = documents[start_index]
+        log.info(f"BATCH: {len(batch)}.")
+        embeddings.extend(model.encode(batch, show_progress_bar=True))
+
+    print(len(embeddings))
+    return embeddings
+
+
+def setup_nltk(temp_dir):
+    """
+    Set up NLTK by creating a temporary directory for NLTK data and downloading required resources.
+    """
+    nltk_data_path = temp_dir / "nltk_data"
+    nltk_data_path.mkdir(exist_ok=True)
+    nltk.data.path.append(str(nltk_data_path))
+
+    nltk.download("stopwords", download_dir=str(nltk_data_path))
+    nltk.download("wordnet", download_dir=str(nltk_data_path))
+
+
+def run(job_input: IJobInput):
+    log.info(f"Starting job step {__name__}")
+
+    data_job_dir = pathlib.Path(job_input.get_job_directory())
+    input_json = data_job_dir / DOCUMENTS_JSON_FILE_LOCATION
+    output_embeddings = data_job_dir / EMBEDDINGS_PKL_FILE_LOCATION
+
+    temp_dir = job_input.get_temporary_write_directory()
+    setup_nltk(temp_dir)
+
+    cleaned_documents = load_and_clean_documents(input_json)
+    if cleaned_documents:
+        log.info(
+            f"{len(cleaned_documents)} documents loaded and cleaned for embedding."
+        )
+        embeddings = embed_documents_in_batches(cleaned_documents)
+        with open(output_embeddings, "wb") as file:
+            import pickle
+
+            pickle.dump(embeddings, file)
+        log.info(f"Embeddings saved to {output_embeddings}")
diff --git a/examples/embed-ingest-job-example/30_create_schema.sql b/examples/embed-ingest-job-example/30_create_schema.sql
@@ -0,0 +1,19 @@
+DROP TABLE IF EXISTS public.vdk_confluence_doc_embeddings_example CASCADE;
+DROP TABLE IF EXISTS public.vdk_confluence_doc_metadata_example CASCADE;
+
+
+CREATE TABLE IF NOT EXISTS public.vdk_confluence_doc_embeddings_example
+(
+    id SERIAL PRIMARY KEY,
+    embedding public.vector
+);
+
+CREATE TABLE IF NOT EXISTS public.vdk_confluence_doc_metadata_example
+(
+    id INTEGER PRIMARY KEY,
+    title TEXT,
+    source TEXT,
+    data TEXT,
+    deleted BOOLEAN,
+    CONSTRAINT fk_metadata_embeddings FOREIGN KEY (id) REFERENCES public.vdk_confluence_doc_embeddings_example(id)
+);
diff --git a/examples/embed-ingest-job-example/40_ingest_embeddings.py b/examples/embed-ingest-job-example/40_ingest_embeddings.py
@@ -0,0 +1,54 @@
+# Copyright 2021-2024 VMware, Inc.
+# SPDX-License-Identifier: Apache-2.0
+import json
+import logging
+import pathlib
+import pickle
+
+import numpy as np
+from config import DOCUMENTS_JSON_FILE_LOCATION
+from config import EMBEDDINGS_PKL_FILE_LOCATION
+from vdk.api.job_input import IJobInput
+
+log = logging.getLogger(__name__)
+
+
+def run(job_input: IJobInput):
+    log.info(f"Starting job step {__name__}")
+
+    data_job_dir = pathlib.Path(job_input.get_job_directory())
+    input_embeddings_path = data_job_dir / EMBEDDINGS_PKL_FILE_LOCATION
+    input_documents_path = data_job_dir / DOCUMENTS_JSON_FILE_LOCATION
+
+    with open(input_embeddings_path, "rb") as file:
+        embeddings = pickle.load(file)
+    with open(input_documents_path) as file:
+        documents = json.load(file)
+
+    print(len(documents), len(embeddings))
+
+    for i, embedding in enumerate(embeddings):
+        embedding_list = (
+            embedding.tolist() if isinstance(embedding, np.ndarray) else embedding
+        )
+        embedding_payload = {
+            "id": documents[i]["metadata"]["id"],
+            "embedding": embedding_list,
+        }
+        job_input.send_object_for_ingestion(
+            payload=embedding_payload,
+            destination_table="vdk_confluence_doc_embeddings_example",
+        )
+
+    for document in documents:
+        metadata_payload = {
+            "id": document["metadata"]["id"],
+            "title": document["metadata"]["title"],
+            "data": document["data"],
+            "source": document["metadata"]["source"],
+            "deleted": document["metadata"]["deleted"],
+        }
+        job_input.send_object_for_ingestion(
+            payload=metadata_payload,
+            destination_table="vdk_confluence_doc_metadata_example",
+        )
diff --git a/examples/embed-ingest-job-example/README.md b/examples/embed-ingest-job-example/README.md
@@ -0,0 +1,27 @@
+# Embed And Ingest Data Job Example
+
+The following Versatile Data Kit example allows you to embed your Confluence JSON data
+and ingest it into Postgres instance with pgvector.
+
+# Create embeddings for the data
+The fetched data from the previous step is read, cleaned and embedded using the
+[all-mpnet-base-v2](https://huggingface.co/sentence-transformers/all-mpnet-base-v2) HuggingFace SentenceTransformer Embedding model.
+
+To open the output embeddings pickle file, use:
+
+```python
+import pandas as pd
+
+obj = pd.read_pickle(r'embeddings_example.pkl')
+```
+
+# Ingest into Postgres
+
+In order to connect to the database, we use [vdk-postgres](https://github.com/vmware/versatile-data-kit/tree/main/projects/vdk-plugins/vdk-postgres).
+You should set the relevant postgres parameters for your instance in the config.ini file.
+
+# Run the example
+To run the data job locally:
+```bash
+vdk run embed-ingest-job-example
+```
diff --git a/examples/embed-ingest-job-example/config.ini b/examples/embed-ingest-job-example/config.ini
@@ -0,0 +1,12 @@
+[owner]
+team = supercollider
+
+[vdk]
+db_default_type=POSTGRES
+ingest_method_default=POSTGRES
+postgres_dbname=
+postgres_dsn=
+postgres_host=
+postgres_password=
+postgres_user=
+ingester_wait_to_finish_after_every_send=True
diff --git a/examples/embed-ingest-job-example/config.py b/examples/embed-ingest-job-example/config.py
@@ -0,0 +1,5 @@
+# Copyright 2021-2024 VMware, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+DOCUMENTS_JSON_FILE_LOCATION = "documents_example.json"
+EMBEDDINGS_PKL_FILE_LOCATION = "embeddings_example.pkl"
diff --git a/examples/embed-ingest-job-example/documents_example.json b/examples/embed-ingest-job-example/documents_example.json
@@ -0,0 +1,47 @@
+[
+  {
+    "metadata": {
+      "title": "Getting Started",
+      "id": "123213312",
+      "source": "https://github.com/vmware/versatile-data-kit/wiki/Getting-Started",
+      "deleted": false
+    },
+    "data": "VDK Getting Started guide"
+  },
+  {
+    "metadata": {
+      "title": "VDK Wiki",
+      "id": "747124724",
+      "source": "https://github.com/vmware/versatile-data-kit/wiki",
+      "deleted": false
+    },
+    "data": "VDK Wiki"
+  },
+  {
+    "metadata": {
+      "title": "VDK Issues",
+      "id": "721295269",
+      "source": "https://github.com/vmware/versatile-data-kit/issues",
+      "deleted": false
+    },
+    "data": "VDK Issues"
+  },
+  {
+    "metadata": {
+      "title": "VDK PRs",
+      "id": "1323122133",
+      "source": "https://github.com/vmware/versatile-data-kit/pulls",
+      "deleted": false
+    },
+    "data": "VDK Pull Requests"
+  },
+  {
+    "metadata": {
+      "title": "VDK Main Page",
+      "id": "312343243",
+      "source": "https://github.com/vmware/versatile-data-kit/tree/main",
+      "deleted": false
+    },
+    "data": "VDK: One framework to develop, deploy and operate data workflows with Python and SQL."
+  }
+]
diff --git a/examples/embed-ingest-job-example/requirements.txt b/examples/embed-ingest-job-example/requirements.txt
@@ -0,0 +1,4 @@
+nltk
+numpy
+sentence-transformers
+vdk-postgres