diff --git a/BUILD.yaml b/BUILD.yaml index 5df92fee7..333a6af04 100644 --- a/BUILD.yaml +++ b/BUILD.yaml @@ -631,3 +631,16 @@ tests_path: tests/fintech_quant/ command: bash tests.sh timeout_in_sec: 1800 + +# owner: @daniel.arrizza +- name: ecommerce_end_to_end + dir: templates/ecommerce_end_to_end + cluster_env: + image_uri: anyscale/ray:2.55.1-py312 + compute_config: + AWS: configs/ecommerce_end_to_end/aws.yaml + GCP: configs/ecommerce_end_to_end/gce.yaml + test: + tests_path: tests/ecommerce_end_to_end/ + command: bash tests.sh + timeout_in_sec: 1800 diff --git a/configs/ecommerce_end_to_end/aws.yaml b/configs/ecommerce_end_to_end/aws.yaml new file mode 100644 index 000000000..9dae0258a --- /dev/null +++ b/configs/ecommerce_end_to_end/aws.yaml @@ -0,0 +1,7 @@ +head_node_type: + name: head + instance_type: m5.2xlarge + +worker_node_types: [] + +auto_select_worker_config: true diff --git a/configs/ecommerce_end_to_end/gce.yaml b/configs/ecommerce_end_to_end/gce.yaml new file mode 100644 index 000000000..55a925465 --- /dev/null +++ b/configs/ecommerce_end_to_end/gce.yaml @@ -0,0 +1,7 @@ +head_node_type: + name: head + instance_type: n2-standard-8 + +worker_node_types: [] + +auto_select_worker_config: true diff --git a/templates/ecommerce_end_to_end/.anyscaleignore b/templates/ecommerce_end_to_end/.anyscaleignore new file mode 100644 index 000000000..30c59efc9 --- /dev/null +++ b/templates/ecommerce_end_to_end/.anyscaleignore @@ -0,0 +1,18 @@ +# This file is used to exclude files from Anyscale Workspaces snapshots. +# Use this to prevent large or unnecessary files from being included in your snapshots, +# which helps reduce snapshot size and creation time. See documentation for more details: +# https://docs.anyscale.com/platform/workspaces/workspaces-files/#excluding-files-with-anyscaleignore +# +# Syntax examples: +# *.txt # Ignore files with a .txt extension at the same level as `.anyscaleignore`. +# **/*.txt # Ignore files with a .txt extension in ANY directory. +# folder/ # Ignore all files under "folder/". The slash at the end is optional. +# folder/*.txt # Ignore files with a .txt extension under "folder/". +# path/to/filename.py # Ignore a specific file by providing its relative path. +# file_[1,2].txt # Ignore file_1.txt and file_2.txt. + +# Exclude Python virtual environments (.venv/) from snapshots. Virtual environments contain +# all installed Python dependencies, which can be multiple gigabytes in size. These directories +# are typically recreatable from requirements files and don't need to be included in snapshots. +# The ** pattern ensures all .venv directories are excluded regardless of location in your project. +**/.venv/ diff --git a/templates/ecommerce_end_to_end/README.md b/templates/ecommerce_end_to_end/README.md new file mode 100644 index 000000000..118988dbb --- /dev/null +++ b/templates/ecommerce_end_to_end/README.md @@ -0,0 +1,5 @@ +# E-Commerce Recommendation System Demo + +An end-to-end recommendation system using Ray Data, Ray Train, and Ray Serve on Anyscale. + +The main notebook is [notebook.ipynb](./notebook.ipynb) \ No newline at end of file diff --git a/templates/ecommerce_end_to_end/client.py b/templates/ecommerce_end_to_end/client.py new file mode 100644 index 000000000..bd7abc532 --- /dev/null +++ b/templates/ecommerce_end_to_end/client.py @@ -0,0 +1,83 @@ +""" +Quick test client for the Ray Serve recommendation endpoint. + +Usage: + # Start the service first: + # serve run serve_app:app + # + # Then in another terminal: + python client.py +""" + +import base64 +import io +import json +import sys + +import requests +from PIL import Image + +SERVE_URL = "http://localhost:8000" + + +def encode_image(path_or_pil) -> str: + """Return base64-encoded JPEG string from a path or PIL Image.""" + if isinstance(path_or_pil, str): + with open(path_or_pil, "rb") as f: + raw = f.read() + else: + buf = io.BytesIO() + path_or_pil.save(buf, format="JPEG") + raw = buf.getvalue() + return base64.b64encode(raw).decode("utf-8") + + +def test_health(): + print("=== Health check ===") + resp = requests.get(f"{SERVE_URL}/health", timeout=5) + print(f"Status: {resp.status_code} Body: {resp.json()}") + print() + + +def test_recommend_demo(): + """Send a real product image and call /recommend.""" + from utils import get_product_image, PRODUCTS + + print("=== /recommend — demo product image ===") + + # Use the first product (Wireless Headphones) as query image + product = PRODUCTS[0] + img_arr = get_product_image(product) + img_pil = Image.fromarray(img_arr) + + payload = {"image_base64": encode_image(img_pil)} + + try: + resp = requests.post( + f"{SERVE_URL}/recommend", + json=payload, + timeout=120, + ) + resp.raise_for_status() + result = resp.json() + except requests.exceptions.ConnectionError: + print( + "ERROR: Could not connect. " + "Make sure `serve run serve_app:app` is running first." + ) + sys.exit(1) + + print(f"Caption: {result['caption']!r}") + print(f"\nTop {len(result['recommendations'])} recommendations:") + for r in result["recommendations"]: + print( + f" {r['rank']}. [{r['category']:18s}] {r['name']:35s} " + f"sim={r['similarity']:.3f}" + ) + print() + + +if __name__ == "__main__": + test_health() + test_recommend_demo() + print("✅ All tests passed!") diff --git a/templates/ecommerce_end_to_end/data/demo_images/Air_Purifier.jpg b/templates/ecommerce_end_to_end/data/demo_images/Air_Purifier.jpg new file mode 100644 index 000000000..d4f65357b Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Air_Purifier.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Bamboo_Cutting_Board.jpg b/templates/ecommerce_end_to_end/data/demo_images/Bamboo_Cutting_Board.jpg new file mode 100644 index 000000000..3a9580440 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Bamboo_Cutting_Board.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Bluetooth_Speaker.jpg b/templates/ecommerce_end_to_end/data/demo_images/Bluetooth_Speaker.jpg new file mode 100644 index 000000000..ee541c71d Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Bluetooth_Speaker.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Cast_Iron_Skillet.jpg b/templates/ecommerce_end_to_end/data/demo_images/Cast_Iron_Skillet.jpg new file mode 100644 index 000000000..5af51fe6c Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Cast_Iron_Skillet.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Clean_Code.jpg b/templates/ecommerce_end_to_end/data/demo_images/Clean_Code.jpg new file mode 100644 index 000000000..61a3f8f85 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Clean_Code.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Cotton_T-Shirt_3-Pack.jpg b/templates/ecommerce_end_to_end/data/demo_images/Cotton_T-Shirt_3-Pack.jpg new file mode 100644 index 000000000..a5980b688 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Cotton_T-Shirt_3-Pack.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Deep_Learning_Book.jpg b/templates/ecommerce_end_to_end/data/demo_images/Deep_Learning_Book.jpg new file mode 100644 index 000000000..602c2f88f Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Deep_Learning_Book.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Designing_Data-Intensive_Apps.jpg b/templates/ecommerce_end_to_end/data/demo_images/Designing_Data-Intensive_Apps.jpg new file mode 100644 index 000000000..61a3f8f85 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Designing_Data-Intensive_Apps.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Distributed_Systems.jpg b/templates/ecommerce_end_to_end/data/demo_images/Distributed_Systems.jpg new file mode 100644 index 000000000..61a3f8f85 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Distributed_Systems.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Dumbbell_Set.jpg b/templates/ecommerce_end_to_end/data/demo_images/Dumbbell_Set.jpg new file mode 100644 index 000000000..898b4a72f Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Dumbbell_Set.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Foam_Roller.jpg b/templates/ecommerce_end_to_end/data/demo_images/Foam_Roller.jpg new file mode 100644 index 000000000..0759e1800 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Foam_Roller.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/French_Press_Coffee.jpg b/templates/ecommerce_end_to_end/data/demo_images/French_Press_Coffee.jpg new file mode 100644 index 000000000..12fb0ff84 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/French_Press_Coffee.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Hiking_Boots.jpg b/templates/ecommerce_end_to_end/data/demo_images/Hiking_Boots.jpg new file mode 100644 index 000000000..72c70c1e8 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Hiking_Boots.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Jump_Rope.jpg b/templates/ecommerce_end_to_end/data/demo_images/Jump_Rope.jpg new file mode 100644 index 000000000..b8012ab92 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Jump_Rope.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Laptop_Stand.jpg b/templates/ecommerce_end_to_end/data/demo_images/Laptop_Stand.jpg new file mode 100644 index 000000000..6bb27f9d3 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Laptop_Stand.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Leather_Belt.jpg b/templates/ecommerce_end_to_end/data/demo_images/Leather_Belt.jpg new file mode 100644 index 000000000..7c5416ae7 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Leather_Belt.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Mechanical_Keyboard.jpg b/templates/ecommerce_end_to_end/data/demo_images/Mechanical_Keyboard.jpg new file mode 100644 index 000000000..d540641b0 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Mechanical_Keyboard.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Merino_Wool_Sweater.jpg b/templates/ecommerce_end_to_end/data/demo_images/Merino_Wool_Sweater.jpg new file mode 100644 index 000000000..5e0123535 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Merino_Wool_Sweater.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Portable_SSD_1TB.jpg b/templates/ecommerce_end_to_end/data/demo_images/Portable_SSD_1TB.jpg new file mode 100644 index 000000000..23a568247 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Portable_SSD_1TB.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Puffer_Vest.jpg b/templates/ecommerce_end_to_end/data/demo_images/Puffer_Vest.jpg new file mode 100644 index 000000000..df77f2627 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Puffer_Vest.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Python_Cookbook.jpg b/templates/ecommerce_end_to_end/data/demo_images/Python_Cookbook.jpg new file mode 100644 index 000000000..61a3f8f85 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Python_Cookbook.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Resistance_Bands_Set.jpg b/templates/ecommerce_end_to_end/data/demo_images/Resistance_Bands_Set.jpg new file mode 100644 index 000000000..dffa6180a Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Resistance_Bands_Set.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Running_Shoes.jpg b/templates/ecommerce_end_to_end/data/demo_images/Running_Shoes.jpg new file mode 100644 index 000000000..b3ebff188 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Running_Shoes.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Scented_Soy_Candle.jpg b/templates/ecommerce_end_to_end/data/demo_images/Scented_Soy_Candle.jpg new file mode 100644 index 000000000..ecb1d48bc Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Scented_Soy_Candle.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Slim-Fit_Chinos.jpg b/templates/ecommerce_end_to_end/data/demo_images/Slim-Fit_Chinos.jpg new file mode 100644 index 000000000..73b899114 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Slim-Fit_Chinos.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Smart_Watch.jpg b/templates/ecommerce_end_to_end/data/demo_images/Smart_Watch.jpg new file mode 100644 index 000000000..431d8419f Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Smart_Watch.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Succulent_Set_6-Pack.jpg b/templates/ecommerce_end_to_end/data/demo_images/Succulent_Set_6-Pack.jpg new file mode 100644 index 000000000..43ba673c9 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Succulent_Set_6-Pack.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/The_Algorithm_Design.jpg b/templates/ecommerce_end_to_end/data/demo_images/The_Algorithm_Design.jpg new file mode 100644 index 000000000..61a3f8f85 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/The_Algorithm_Design.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/USB-C_Hub.jpg b/templates/ecommerce_end_to_end/data/demo_images/USB-C_Hub.jpg new file mode 100644 index 000000000..b2851ba4c Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/USB-C_Hub.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Water_Bottle_32oz.jpg b/templates/ecommerce_end_to_end/data/demo_images/Water_Bottle_32oz.jpg new file mode 100644 index 000000000..48446fec6 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Water_Bottle_32oz.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Waterproof_Jacket.jpg b/templates/ecommerce_end_to_end/data/demo_images/Waterproof_Jacket.jpg new file mode 100644 index 000000000..77e6159cb Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Waterproof_Jacket.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Webcam_1080p.jpg b/templates/ecommerce_end_to_end/data/demo_images/Webcam_1080p.jpg new file mode 100644 index 000000000..3e1edaa5b Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Webcam_1080p.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Wireless_Headphones.jpg b/templates/ecommerce_end_to_end/data/demo_images/Wireless_Headphones.jpg new file mode 100644 index 000000000..06c237844 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Wireless_Headphones.jpg differ diff --git a/templates/ecommerce_end_to_end/data/demo_images/Yoga_Mat.jpg b/templates/ecommerce_end_to_end/data/demo_images/Yoga_Mat.jpg new file mode 100644 index 000000000..b6eda7dc6 Binary files /dev/null and b/templates/ecommerce_end_to_end/data/demo_images/Yoga_Mat.jpg differ diff --git a/templates/ecommerce_end_to_end/notebook.ipynb b/templates/ecommerce_end_to_end/notebook.ipynb new file mode 100644 index 000000000..d714230ec --- /dev/null +++ b/templates/ecommerce_end_to_end/notebook.ipynb @@ -0,0 +1,480 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "cell-0", + "metadata": {}, + "source": [ + "# E-Commerce Recommendation System\n", + "\n", + "End-to-end demo using **Ray Data → Ray Train → Ray Data → Ray Serve** on Anyscale.\n", + "\n", + "```\n", + "┌──────────────────────────────────────────────────────────────────────┐\n", + "│ Product Recommender │\n", + "│ │\n", + "│ ▼ Stage 1 — Ray Data │\n", + "│ Preprocess images (resize, normalise) │\n", + "│ Clean text (descriptions + category) │\n", + "│ │ │\n", + "│ ▼ Stage 2 — Ray Train (TorchTrainer) │\n", + "│ Fine-tune all-MiniLM-L6-v2 with contrastive loss │\n", + "│ Save fine-tuned model checkpoint │\n", + "│ │ │\n", + "│ ▼ Stage 3 — Ray Data (Batch Embedding) │\n", + "│ Embed entire product catalog with fine-tuned model │\n", + "│ Save embeddings + metadata for serving │\n", + "│ │ │\n", + "│ ▼ Stage 4 — Ray Serve │\n", + "│ POST /recommend {image_base64} │\n", + "│ ImageToText (BLIP-base) → caption │\n", + "│ ProductRecommender (MiniLM) → top-5 products │\n", + "│ │ │\n", + "│ ▼ Streamlit UI │\n", + "│ Upload image │\n", + "│ │\n", + "└──────────────────────────────────────────────────────────────────────┘\n", + "```\n", + "\n", + "**Models used**\n", + "| Stage | Model | Size | Device |\n", + "|-------|-------|------|--------|\n", + "| Train | `sentence-transformers/all-MiniLM-L6-v2` | 22 M params | CPU |\n", + "| Embed | fine-tuned MiniLM | 22 M params | CPU |\n", + "| Serve | `Salesforce/blip-image-captioning-base` | 224 M params | CPU |\n", + "| Serve | fine-tuned MiniLM | 22 M params | CPU |\n", + "\n", + "Everything runs on **CPU** and completes in a few minutes." + ] + }, + { + "cell_type": "markdown", + "id": "cell-1", + "metadata": {}, + "source": [ + "## 0 — Setup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-2", + "metadata": {}, + "outputs": [], + "source": [ + "# Install dependencies (skip if already installed)\n", + "!pip install -q -r setup/requirements.txt" + ] + }, + { + "cell_type": "markdown", + "id": "cell-4", + "metadata": {}, + "source": [ + "### Peek at the product images" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-5", + "metadata": {}, + "outputs": [], + "source": [ + "from utils.viz import plot_category_samples\n", + "import json\n", + "import os\n", + "import sys\n", + "import numpy as np\n", + "from pathlib import Path\n", + "from PIL import Image\n", + "# Make sure the project root is on the path\n", + "sys.path.insert(0, str(Path().resolve()))\n", + "\n", + "from utils import PRODUCTS, CATEGORIES, get_product_image\n", + "\n", + "plot_category_samples(PRODUCTS, CATEGORIES, get_product_image)" + ] + }, + { + "cell_type": "markdown", + "id": "cell-6", + "metadata": {}, + "source": [ + "---\n", + "## Stage 1 — Catalog Generation\n", + "\n", + "Here we expand the base 30-product catalog to 1,000 products and add `text_clean` (strip / lowercase / collapse whitespace on `training_text`) directly in memory — no Parquet intermediate needed." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-7", + "metadata": {}, + "outputs": [], + "source": [ + "import ray\n", + "from utils import (\n", + " PRODUCTS,\n", + " attach_clean_text,\n", + " expand_catalog,\n", + " generate_catalog,\n", + " init_ray,\n", + ")\n", + "\n", + "init_ray()\n", + "\n", + "raw_dir = os.path.abspath(\"data/raw\")\n", + "\n", + "products = expand_catalog(PRODUCTS, target_size=1000)\n", + "print(f\"Catalog size : {len(products)} products\")\n", + "print(f\"Categories : {sorted(set(p['category'] for p in products))}\")\n", + "\n", + "catalog_records = attach_clean_text(generate_catalog(products=products, output_dir=raw_dir))\n", + "print(f\"Catalog records : {len(catalog_records)}\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "cell-9", + "metadata": {}, + "source": [ + "---\n", + "## Stage 2 — Ray Train: Embedding Fine-Tuning\n", + "\n", + "The starter model (`all-MiniLM-L6-v2`) learned from general web text, not from our product categories. That makes it easy for a fuzzy query to land near the wrong kind of item in the embedding space.\n", + "\n", + "**Fine-tuning** trains the model with **contrastive loss**: same-category products move closer, different-category products move apart, so search and recommendations follow category boundaries better.\n", + "\n", + "We train on **100 products (20 per category)** so each category is represented and runs stay short. The t-SNE plots below show embeddings before vs after fine-tuning." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-10a", + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "from ray.train import CheckpointConfig, FailureConfig, RunConfig, ScalingConfig\n", + "from ray.train.torch import TorchTrainer\n", + "from utils.training import train_loop_per_worker\n", + "from utils import sample_per_category, resolve_artifact_paths\n", + "\n", + "TRAIN_PER_CATEGORY = 20\n", + "\n", + "# set up some paths\n", + "paths = resolve_artifact_paths()\n", + "MODEL_OUTPUT_DIR = paths[\"model_dir\"]\n", + "EMBEDDINGS_PATH = paths[\"embeddings_path\"]\n", + "METADATA_PATH = paths[\"metadata_path\"]\n", + "TRAIN_RESULT_DIR = paths[\"train_result_dir\"]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "71a65499", + "metadata": {}, + "outputs": [], + "source": [ + "# TorchTrainer: 1 worker, contrastive loss pulls same-category pairs together\n", + "\n", + "trainer = TorchTrainer(\n", + " train_loop_per_worker=train_loop_per_worker,\n", + " train_loop_config={\n", + " \"base_model\": \"sentence-transformers/all-MiniLM-L6-v2\",\n", + " \"epochs\": 2,\n", + " \"batch_size\": 8,\n", + " \"lr\": 2e-5,\n", + " \"seed\": 42,\n", + " \"records\": sample_per_category(\n", + " catalog_records, n_per_category=TRAIN_PER_CATEGORY, seed=42\n", + " ),\n", + " },\n", + " scaling_config=ScalingConfig(num_workers=1, use_gpu=torch.cuda.is_available()),\n", + " run_config=RunConfig(\n", + " name=\"ecomm_embedding_finetune\",\n", + " storage_path=os.path.abspath(TRAIN_RESULT_DIR),\n", + " checkpoint_config=CheckpointConfig(\n", + " num_to_keep=2,\n", + " checkpoint_score_attribute=\"train_loss\",\n", + " checkpoint_score_order=\"min\",\n", + " ),\n", + " failure_config=FailureConfig(max_failures=1),\n", + " ),\n", + ")\n", + "print(\"TorchTrainer configured\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a52533d2", + "metadata": {}, + "outputs": [], + "source": [ + "# Run the fine-tune and persist the best checkpoint in a ready-to-serve layout.\n", + "from utils.training import save_best_sentence_transformer\n", + "\n", + "result = trainer.fit()\n", + "print(result.metrics_dataframe[[\"epoch\", \"train_loss\"]].to_string(index=False))\n", + "\n", + "save_best_sentence_transformer(result, MODEL_OUTPUT_DIR)\n", + "print(f\"Model saved → {MODEL_OUTPUT_DIR}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "mfak14ffr38", + "metadata": {}, + "outputs": [], + "source": [ + "from utils.viz import plot_training_loss\n", + "\n", + "plot_training_loss(result.metrics_dataframe)" + ] + }, + { + "cell_type": "markdown", + "id": "stage3-md", + "metadata": {}, + "source": [ + "---\n", + "## Stage 3 — Ray Data: Batch Embedding\n", + "\n", + "**The problem this solves:** embedding 1,000 products serially with a model loaded once per process is slow. Worse, at 1M products you can't fit the whole catalog in memory on one machine.\n", + "\n", + "Ray Data solves both: each actor loads the fine-tuned model **once**, then processes a steady stream of batches. Workers run in parallel, and Ray's streaming execution means only a slice of the catalog is in memory at any time. The same code scales from 1,000 to 100M products by adjusting the `ActorPoolStrategy` size." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19ef0810", + "metadata": {}, + "outputs": [], + "source": [ + "# Actor pool: each worker loads the fine-tuned model once and encodes batches in parallel\n", + "from pathlib import Path\n", + "from utils.embedding import ProductEmbedder\n", + "import ray\n", + "from utils import init_ray\n", + "\n", + "init_ray()\n", + "\n", + "assert Path(MODEL_OUTPUT_DIR).exists() and any(Path(MODEL_OUTPUT_DIR).iterdir()), \\\n", + " f\"Model not found at {MODEL_OUTPUT_DIR} — run Stage 2 first\"\n", + "\n", + "ds = ray.data.from_items(catalog_records).select_columns(\n", + " [\"product_id\", \"name\", \"category\", \"text_clean\"]\n", + ")\n", + "\n", + "rows = ds.map_batches(\n", + " ProductEmbedder,\n", + " fn_constructor_kwargs={\"model_dir\": MODEL_OUTPUT_DIR},\n", + " batch_size=8,\n", + " num_cpus=1,\n", + " compute=ray.data.ActorPoolStrategy(size=2),\n", + " batch_format=\"numpy\",\n", + ").take_all()\n", + "print(f\"Embedded {len(rows)} products\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10e45061", + "metadata": {}, + "outputs": [], + "source": [ + "# Persist vectors and metadata as two separate files.\n", + "#\n", + "# The dense (N, 384) matrix goes to .npy for fast mmap-style loads at serve\n", + "# time; product name/category/id go to a small JSON sidecar so the serving\n", + "# layer never has to pull full product text into memory at query time.\n", + "from utils.embedding import save_embeddings_and_metadata\n", + "\n", + "embeddings, metadata = save_embeddings_and_metadata(rows, EMBEDDINGS_PATH, METADATA_PATH)\n", + "print(f\"Embeddings : {embeddings.shape} → {EMBEDDINGS_PATH}\")\n", + "print(f\"Metadata : {len(metadata)} products → {METADATA_PATH}\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-11", + "metadata": {}, + "outputs": [], + "source": [ + "from utils.viz import plot_similarity_heatmap\n", + "\n", + "embeddings = np.load(EMBEDDINGS_PATH)\n", + "with open(METADATA_PATH) as f:\n", + " metadata = json.load(f)\n", + "\n", + "print(f\"Embeddings shape : {embeddings.shape}\")\n", + "print(f\"Products indexed : {len(metadata)}\")\n", + "\n", + "plot_similarity_heatmap(embeddings, [m[\"name\"] for m in metadata])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4y3rm2m16vy", + "metadata": {}, + "outputs": [], + "source": [ + "import importlib\n", + "\n", + "import utils.viz as _viz\n", + "\n", + "importlib.reload(_viz)\n", + "from utils.viz import plot_tsne_comparison\n", + "\n", + "base_embs, ft_embs, aligned_meta = plot_tsne_comparison(\n", + " metadata, embeddings, catalog_records=catalog_records\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "zf5cnnozpe", + "metadata": {}, + "outputs": [], + "source": [ + "# Numeric check: did fine-tuning actually make the geometry category-aware?\n", + "#\n", + "# Simple \"top-k neighbors share the same category\" scores read ~100% on this\n", + "# demo because product names already sound like their category — so the base\n", + "# model looks strong before we even train. The two metrics below still move\n", + "# after fine-tuning and tell us the embedding space really did reshape.\n", + "from utils.viz import print_embedding_quality_report\n", + "\n", + "print_embedding_quality_report(base_embs, ft_embs, aligned_meta)" + ] + }, + { + "cell_type": "markdown", + "id": "cell-12", + "metadata": {}, + "source": [ + "---\n", + "## Stage 4 — Ray Serve: Online Recommendation API\n", + "\n", + "**The problem this solves:** BLIP (224M params) and MiniLM (22M params) have very different latency, throughput, and resource profiles. Bundling them into one process means the slower model throttles the faster one.\n", + "\n", + "Ray Serve deploys them as **independent deployments** that auto-scale separately. `ImageToText` can scale to 4 replicas for a spike in image uploads without adding MiniLM replicas you don't need. Both can be upgraded or rolled back independently, and the composing `RecommendationService` wires them together asynchronously behind a single `/recommend` endpoint." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-13", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import ray\n", + "from ray import serve\n", + "from utils import init_ray\n", + "\n", + "init_ray()\n", + "\n", + "os.environ[\"EMBEDDING_MODEL_DIR\"] = MODEL_OUTPUT_DIR\n", + "os.environ[\"EMBEDDINGS_PATH\"] = EMBEDDINGS_PATH\n", + "os.environ[\"METADATA_PATH\"] = METADATA_PATH\n", + "\n", + "import importlib, serve_app\n", + "importlib.reload(serve_app)\n", + "from serve_app import app\n", + "\n", + "serve.run(app)\n", + "print(\"Service running at http://localhost:8000\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-14", + "metadata": {}, + "outputs": [], + "source": [ + "# Smoke-test the running /recommend endpoint with a known product image.\n", + "#\n", + "# `encode_image_base64` handles JPEG compression + base64 in one call, and\n", + "# `post_recommend` POSTs the payload shape the service expects.\n", + "from utils import PRODUCTS, get_product_image, post_recommend\n", + "\n", + "product = PRODUCTS[0] # Wireless Headphones — should match other Electronics\n", + "img_arr = get_product_image(product)\n", + "result = post_recommend(img_arr)\n", + "\n", + "print(f\"Caption : {result['caption']!r}\")\n", + "print(f\"\\nTop-{len(result['recommendations'])} recommendations:\")\n", + "for r in result[\"recommendations\"]:\n", + " print(\n", + " f\" {r['rank']}. [{r['category']:18s}] {r['name']:35s} sim={r['similarity']:.3f}\"\n", + " )\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8vktgpfk9d", + "metadata": {}, + "outputs": [], + "source": [ + "# One image per category — does the pipeline generalise beyond headphones?\n", + "from utils import post_recommend\n", + "\n", + "print(f\"{'Category':<16} {'Caption':<45} {'Top recommendation'}\")\n", + "print(\"-\" * 85)\n", + "for cat in CATEGORIES:\n", + " prod = next(p for p in PRODUCTS if p[\"category\"] == cat)\n", + " r = post_recommend(get_product_image(prod))\n", + " top = r[\"recommendations\"][0]\n", + " print(\n", + " f\"{cat:<16} {r['caption'][:44]:<45} {top['name'][:20]} ({top['similarity']:.2f})\"\n", + " )\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-16", + "metadata": {}, + "outputs": [], + "source": [ + "# Teardown (optional — keep the service running if you want to use the Streamlit app)\n", + "# serve.shutdown()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "base", + "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.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/templates/ecommerce_end_to_end/serve_app.py b/templates/ecommerce_end_to_end/serve_app.py new file mode 100644 index 000000000..dbe138f6b --- /dev/null +++ b/templates/ecommerce_end_to_end/serve_app.py @@ -0,0 +1,356 @@ +""" +E-Commerce Recommendation System — Ray Serve Application +========================================================= + +Stage 3 of 3: Online recommendation endpoint + +Pipeline (multi-model composition) +----------------------------------- + POST /recommend {"image_base64": ""} + │ + ▼ + ImageToText ← Salesforce/blip-image-captioning-base + │ caption text + ▼ + ProductRecommender ← fine-tuned all-MiniLM-L6-v2 + cosine index + │ + ▼ + JSON response {"caption": "...", "recommendations": [...]} + +Run locally (development) +-------------------------- + serve run serve_app:app + # or: + python serve_app.py + +Run as an Anyscale Service +--------------------------- + anyscale service deploy -f service.yaml + +Test +---- + python client.py + +Ray version: 2.x (Ray ≥ 2.20) +Base image: anyscale/ray:2.47.1-slim-py312 (CPU-only) +See https://docs.anyscale.com/reference/base-images for the latest images. + +References +---------- +- Ray Serve Model Composition: https://docs.ray.io/en/latest/serve/model_composition.html +- Ray Serve Key Concepts: https://docs.ray.io/en/latest/serve/key-concepts.html +- Anyscale Services: https://docs.anyscale.com/services/ +""" + +import base64 +import io +import json +import os +from pathlib import Path +from typing import Any + +import numpy as np +from fastapi import FastAPI, HTTPException +from pydantic import BaseModel +from ray import serve +from ray.serve.handle import DeploymentHandle +from starlette.requests import Request + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + +BLIP_MODEL = "Salesforce/blip-image-captioning-base" +_HERE = Path(__file__).parent.resolve() + +# Model resolution order (first match wins): +# 1. Env var EMBEDDING_MODEL_DIR (explicit override) +# 2. Cluster storage fine-tuned model (written by notebook when /mnt/cluster_storage exists) +# 3. Local models/embedding_model (pushed or generated locally) +# 4. HuggingFace model ID — pre-cached in container image (Dockerfile bakes it in) +_CLUSTER_MODEL = Path("/mnt/cluster_storage") / "ecomm_embedding_model" +_LOCAL_MODEL = _HERE / "models/embedding_model" +_default_model_dir = ( + str(_CLUSTER_MODEL) if _CLUSTER_MODEL.exists() + else str(_LOCAL_MODEL) if _LOCAL_MODEL.exists() + else "sentence-transformers/all-MiniLM-L6-v2" +) + +_CLUSTER_EMBEDDINGS = Path("/mnt/cluster_storage") / "ecomm_product_embeddings.npy" +_LOCAL_EMBEDDINGS = _HERE / "models/product_embeddings.npy" +_default_embeddings = ( + str(_CLUSTER_EMBEDDINGS) if _CLUSTER_EMBEDDINGS.exists() + else str(_LOCAL_EMBEDDINGS) +) + +_CLUSTER_METADATA = Path("/mnt/cluster_storage") / "ecomm_product_metadata.json" +_LOCAL_METADATA = _HERE / "models/product_metadata.json" +_default_metadata = ( + str(_CLUSTER_METADATA) if _CLUSTER_METADATA.exists() + else str(_LOCAL_METADATA) +) + +EMBEDDING_MODEL_DIR = os.environ.get("EMBEDDING_MODEL_DIR", _default_model_dir) +EMBEDDINGS_PATH = os.environ.get("EMBEDDINGS_PATH", _default_embeddings) +METADATA_PATH = os.environ.get("METADATA_PATH", _default_metadata) + +TOP_K = int(os.environ.get("TOP_K", "5")) + +# --------------------------------------------------------------------------- +# Pydantic schema +# --------------------------------------------------------------------------- + + +class RecommendRequest(BaseModel): + image_base64: str # standard base64-encoded image bytes + + +# --------------------------------------------------------------------------- +# Deployment 1: Image → Caption (BLIP) +# --------------------------------------------------------------------------- + + +@serve.deployment( + num_replicas=1, + ray_actor_options={"num_cpus": 2}, + health_check_period_s=30, + health_check_timeout_s=60, +) +class ImageToText: + """Load BLIP and convert raw image bytes to a natural-language caption.""" + + def __init__(self): + import torch + from transformers import BlipForConditionalGeneration, BlipProcessor + + self.device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"[ImageToText] Loading {BLIP_MODEL} on {self.device} …") + + self.processor = BlipProcessor.from_pretrained(BLIP_MODEL) + self.model = BlipForConditionalGeneration.from_pretrained(BLIP_MODEL) + self.model.to(self.device).eval() + + # Warm-up pass to avoid first-request latency spike + from PIL import Image as PILImage + + dummy = PILImage.new("RGB", (224, 224), color=(128, 128, 128)) + inputs = self.processor(images=dummy, return_tensors="pt").to(self.device) + import torch + + with torch.no_grad(): + self.model.generate(**inputs, max_new_tokens=30) + + print(f"[ImageToText] Ready.") + + def caption(self, image_bytes: bytes) -> str: + """Generate a caption for the given raw image bytes.""" + import torch + from PIL import Image as PILImage + + img = PILImage.open(io.BytesIO(image_bytes)).convert("RGB") + inputs = self.processor(images=img, return_tensors="pt").to(self.device) + + with torch.no_grad(): + out = self.model.generate( + **inputs, + max_new_tokens=40, + num_beams=3, + ) + + caption = self.processor.decode(out[0], skip_special_tokens=True) + return caption.strip() + + def check_health(self): + pass # model loaded in __init__; if it failed, the actor died + + +# --------------------------------------------------------------------------- +# Deployment 2: Text → Top-K products (sentence-transformer + cosine index) +# --------------------------------------------------------------------------- + + +@serve.deployment( + num_replicas=1, + ray_actor_options={"num_cpus": 2}, + health_check_period_s=30, + health_check_timeout_s=60, +) +class ProductRecommender: + """ + Embed a query text with the fine-tuned sentence-transformer, then return + the top-K most similar products via cosine similarity over a pre-computed + embedding matrix. + + Paths are accepted as constructor arguments so Ray Serve workers (which + do not inherit the driver's ``os.environ`` from a notebook) still load the + same model and index files as the process that called ``bind()``. + """ + + def __init__( + self, + embedding_model_dir: str | None = None, + embeddings_path: str | None = None, + metadata_path: str | None = None, + ): + from sentence_transformers import SentenceTransformer + + model_dir = ( + embedding_model_dir + if embedding_model_dir is not None + else EMBEDDING_MODEL_DIR + ) + emb_path = ( + embeddings_path if embeddings_path is not None else EMBEDDINGS_PATH + ) + meta_path = metadata_path if metadata_path is not None else METADATA_PATH + + print(f"[ProductRecommender] Loading embedding model from {model_dir} …") + self.model = SentenceTransformer(model_dir) + + print(f"[ProductRecommender] Loading product index …") + self.embeddings = np.load(emb_path).astype(np.float32) # (N, D) + # L2-normalise for fast cosine via dot product + norms = np.linalg.norm(self.embeddings, axis=1, keepdims=True) + 1e-9 + self.embeddings_norm = self.embeddings / norms + + with open(meta_path) as f: + self.metadata = json.load(f) # list of {product_id, name, category} + + print(f"[ProductRecommender] Ready. {len(self.metadata)} products indexed.") + + def recommend(self, query_text: str, top_k: int = TOP_K) -> list[dict]: + """Return the top-K products most similar to query_text.""" + q_emb = self.model.encode([query_text], convert_to_numpy=True)[0].astype( + np.float32 + ) + q_norm = q_emb / (np.linalg.norm(q_emb) + 1e-9) + + sims = self.embeddings_norm @ q_norm # (N,) + top_idx = np.argsort(sims)[::-1][:top_k] + + results = [] + for rank, idx in enumerate(top_idx, start=1): + m = self.metadata[idx] + results.append( + { + "rank": rank, + "product_id": m["product_id"], + "name": m["name"], + "category": m["category"], + "similarity": float(round(float(sims[idx]), 4)), + } + ) + return results + + def check_health(self): + pass + + +# --------------------------------------------------------------------------- +# Deployment 3: HTTP ingress (orchestrator) +# --------------------------------------------------------------------------- + +_fastapi = FastAPI( + title="E-Commerce Product Recommender", + description="Upload an image → get product recommendations", + version="1.0.0", +) + + +@serve.deployment( + num_replicas=1, + ray_actor_options={"num_cpus": 1}, +) +@serve.ingress(_fastapi) +class RecommendationService: + """ + HTTP ingress that chains ImageToText → ProductRecommender. + + Both sub-deployments are called via async DeploymentHandle so the + orchstrator never blocks. + """ + + def __init__( + self, + image_to_text: DeploymentHandle, + product_recommender: DeploymentHandle, + ): + self.image_to_text = image_to_text + self.product_recommender = product_recommender + print("[RecommendationService] Ready.") + + # ------------------------------------------------------------------ + # GET /health + # ------------------------------------------------------------------ + + @_fastapi.get("/health") + def health(self) -> dict: + return {"status": "ok"} + + # ------------------------------------------------------------------ + # POST /recommend + # ------------------------------------------------------------------ + + @_fastapi.post("/recommend") + async def recommend(self, body: RecommendRequest) -> dict[str, Any]: + """ + Accept a base64-encoded image, caption it with BLIP, then return + the top-K most similar products. + + Request body: + {"image_base64": ""} + + Response: + { + "caption": "a pair of wireless headphones", + "recommendations": [ + {"rank": 1, "product_id": "P0000", "name": "...", + "category": "Electronics", "similarity": 0.92}, + ... + ] + } + """ + # 1. Decode base64 → raw bytes + try: + image_bytes = base64.b64decode(body.image_base64) + except Exception as e: + raise HTTPException(status_code=400, detail=f"Invalid base64 image: {e}") + + # 2. Image → caption (async remote call to ImageToText) + caption: str = await self.image_to_text.caption.remote(image_bytes) + + # 3. Caption → recommendations (async remote call to ProductRecommender) + recommendations: list[dict] = await self.product_recommender.recommend.remote( + caption, TOP_K + ) + + return { + "caption": caption, + "recommendations": recommendations, + } + + +# --------------------------------------------------------------------------- +# Application binding +# --------------------------------------------------------------------------- + +app = RecommendationService.bind( + image_to_text=ImageToText.bind(), + product_recommender=ProductRecommender.bind( + embedding_model_dir=EMBEDDING_MODEL_DIR, + embeddings_path=EMBEDDINGS_PATH, + metadata_path=METADATA_PATH, + ), +) + + +# --------------------------------------------------------------------------- +# Local development entrypoint +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + import ray + + ray.init(ignore_reinit_error=True) + serve.start(http_options={"host": "0.0.0.0", "port": 8000}) + serve.run(app, blocking=True) diff --git a/templates/ecommerce_end_to_end/setup/Dockerfile b/templates/ecommerce_end_to_end/setup/Dockerfile new file mode 100644 index 000000000..416ebf006 --- /dev/null +++ b/templates/ecommerce_end_to_end/setup/Dockerfile @@ -0,0 +1,24 @@ +FROM anyscale/ray:2.55.1-slim-py312-cu129 + +# Hugging Face cache: base image runs builds as `ray`, which cannot write under /root. +ENV HF_HOME=/home/ray/.cache/huggingface +ENV HF_HUB_ENABLE_HF_TRANSFER=1 + +# Inline deps: remote image builds disallow COPY from local paths (see setup/requirements.txt for the list). +RUN python3 -m pip install --no-cache-dir \ + torch \ + transformers \ + sentence-transformers \ + hf_transfer \ + streamlit + +# Bake public weights into the image so Serve cold start skips Hub downloads. +RUN python3 - <<'PY' +from huggingface_hub import snapshot_download + +for repo_id in ( + "Salesforce/blip-image-captioning-base", + "sentence-transformers/all-MiniLM-L6-v2", +): + snapshot_download(repo_id, token=False) +PY diff --git a/templates/ecommerce_end_to_end/setup/requirements.txt b/templates/ecommerce_end_to_end/setup/requirements.txt new file mode 100644 index 000000000..896a3e428 --- /dev/null +++ b/templates/ecommerce_end_to_end/setup/requirements.txt @@ -0,0 +1,8 @@ +torch +transformers +sentence-transformers +hf_transfer +scikit-learn + +# Notebook / UI +streamlit diff --git a/templates/ecommerce_end_to_end/setup/service.yaml b/templates/ecommerce_end_to_end/setup/service.yaml new file mode 100644 index 000000000..6f18e753f --- /dev/null +++ b/templates/ecommerce_end_to_end/setup/service.yaml @@ -0,0 +1,53 @@ +# Anyscale Service Configuration — E-Commerce Recommendation System +# Deploy from project root with: anyscale service deploy -f setup/service.yaml +# Status: anyscale service status --name ecomm-recommender-service + +name: ecomm-recommender-service + +applications: + - name: default + import_path: serve_app:app + route_prefix: / + deployments: + - name: RecommendationService + num_replicas: 1 + max_ongoing_requests: 20 + ray_actor_options: + num_cpus: 1 + + - name: ImageToText + num_replicas: 1 + max_ongoing_requests: 10 + ray_actor_options: + num_cpus: 2 + + - name: ProductRecommender + num_replicas: 1 + max_ongoing_requests: 10 + ray_actor_options: + num_cpus: 2 + +# CPU-only cluster (BLIP-base and MiniLM run fine on CPU for demo traffic) +compute_config: + head_node: + instance_type: m5.2xlarge + worker_nodes: + - instance_type: m5.2xlarge + min_nodes: 0 + max_nodes: 3 + +# Base image — CPU-only, Python 3.12 +# Check https://docs.anyscale.com/reference/base-images for the latest +image_uri: anyscale/ray:2.55.1-slim-py312 + +working_dir: . + +requirements: + - torch + - transformers + - sentence-transformers + - hf_transfer + - pillow + - numpy + - fastapi + - pydantic diff --git a/templates/ecommerce_end_to_end/streamlit_app.py b/templates/ecommerce_end_to_end/streamlit_app.py new file mode 100644 index 000000000..4275dbf53 --- /dev/null +++ b/templates/ecommerce_end_to_end/streamlit_app.py @@ -0,0 +1,163 @@ +""" +E-Commerce Recommendation System — Streamlit Frontend +====================================================== + +Lets you upload a product image and see the top-5 recommended products +returned by the Ray Serve endpoint. + +Run locally (while serve_app is running): + streamlit run streamlit_app.py + +Environment variables: + SERVE_URL Base URL of the Ray Serve service (default: http://localhost:8000) +""" + +import base64 +import io +import os + +import requests +import streamlit as st +from PIL import Image + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + +SERVE_URL = os.environ.get("SERVE_URL", "http://localhost:8000") +RECOMMEND_ENDPOINT = f"{SERVE_URL}/recommend" +HEALTH_ENDPOINT = f"{SERVE_URL}/health" + +st.set_page_config( + page_title="Product Recommender", + page_icon="🛍️", + layout="wide", +) + +# --------------------------------------------------------------------------- +# Sidebar +# --------------------------------------------------------------------------- + +with st.sidebar: + st.subheader("Service Status") + try: + resp = requests.get(HEALTH_ENDPOINT, timeout=3) + if resp.ok: + st.success("✅ Service online") + else: + st.error(f"❌ Service returned {resp.status_code}") + except Exception as e: + st.warning(f"⚠️ Cannot reach service:\n{e}") + + st.markdown("---") + st.markdown(f"**Endpoint:** `{RECOMMEND_ENDPOINT}`") + top_k = st.slider("Number of recommendations", min_value=1, max_value=10, value=5) + + st.markdown("---") + st.markdown( + "**How it works**\n\n" + "1. BLIP captions the image\n" + "2. Fine-tuned MiniLM embeds the caption\n" + "3. Cosine search over the product catalog" + ) + +# --------------------------------------------------------------------------- +# Header +# --------------------------------------------------------------------------- + +st.title("🛍️ E-Commerce Product Recommender") +st.markdown( + "Upload a product image — the system captions it with **BLIP**, embeds the " + "caption with a **fine-tuned MiniLM**, and returns the most similar products." +) + +# --------------------------------------------------------------------------- +# Image upload +# --------------------------------------------------------------------------- + +col_upload, col_preview = st.columns([1, 1], gap="large") + +with col_upload: + st.subheader("Query Image") + uploaded = st.file_uploader( + "Choose an image (JPG / PNG)", + type=["jpg", "jpeg", "png"], + label_visibility="collapsed", + ) + + if uploaded is None: + st.info("No image uploaded — using a demo product image.") + from utils import get_product_image, PRODUCTS + demo_arr = get_product_image(PRODUCTS[0]) # Wireless Headphones + buf = io.BytesIO() + Image.fromarray(demo_arr).save(buf, format="JPEG") + image_bytes = buf.getvalue() + else: + image_bytes = uploaded.read() + + display_image = Image.open(io.BytesIO(image_bytes)).convert("RGB") + st.image(display_image, use_container_width=True) + + run = st.button("🔍 Find Similar Products", type="primary", use_container_width=True) + +# --------------------------------------------------------------------------- +# Call the service +# --------------------------------------------------------------------------- + +with col_preview: + st.subheader("Results") + + if run: + with st.spinner("Analysing image and searching catalog …"): + encoded = base64.b64encode(image_bytes).decode("utf-8") + try: + resp = requests.post( + RECOMMEND_ENDPOINT, json={"image_base64": encoded}, timeout=60 + ) + resp.raise_for_status() + result = resp.json() + except requests.exceptions.ConnectionError: + st.error( + "Could not connect to the service. " + "Make sure `serve run serve_app:app` is running." + ) + st.stop() + except Exception as e: + st.error(f"Request failed: {e}") + if hasattr(e, "response") and e.response is not None: + st.code(e.response.text) + st.stop() + + # Caption — shown prominently so viewers understand what drove the results + caption = result.get("caption", "(no caption)") + st.markdown("**Caption generated by BLIP:**") + st.info(f'"{caption}"') + + st.markdown(f"**Top-{top_k} recommendations:**") + recs = result.get("recommendations", [])[:top_k] + if not recs: + st.warning("No recommendations returned.") + else: + for rec in recs: + sim = rec["similarity"] + st.markdown( + f"**{rec['rank']}. {rec['name']}**   `{rec['category']}`" + ) + st.progress( + sim, + text=f"Similarity: {sim:.1%}", + ) + st.caption(f"Product IDs: {', '.join(r['product_id'] for r in recs)}") + else: + st.markdown("*Upload an image and click **Find Similar Products** to see results.*") + +# --------------------------------------------------------------------------- +# Footer +# --------------------------------------------------------------------------- + +st.markdown("---") +st.caption( + "Built with [Ray Serve](https://docs.ray.io/en/latest/serve/index.html), " + "[BLIP](https://huggingface.co/Salesforce/blip-image-captioning-base), " + "and [Sentence-Transformers](https://www.sbert.net/) on Anyscale." +) diff --git a/templates/ecommerce_end_to_end/utils/__init__.py b/templates/ecommerce_end_to_end/utils/__init__.py new file mode 100644 index 000000000..b93345f31 --- /dev/null +++ b/templates/ecommerce_end_to_end/utils/__init__.py @@ -0,0 +1,444 @@ +""" +Shared utilities for the e-commerce recommendation system demo. +Kept in a single .py file so both notebooks and scripts can import it. +""" + +import io +import json +import os +import random +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +import numpy as np +from PIL import Image, ImageDraw, ImageFont + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +CATEGORIES = ["Electronics", "Clothing", "Books", "Home & Garden", "Sports"] + +CATEGORY_COLORS = { + "Electronics": (70, 130, 180), # steel blue + "Clothing": (220, 90, 120), # rose + "Books": (120, 170, 60), # olive green + "Home & Garden": (200, 150, 50), # golden + "Sports": (80, 180, 140), # teal +} + +PRODUCTS = [ + # Electronics + {"name": "Wireless Headphones", "category": "Electronics", "price": 79.99, "description": "Over-ear noise-cancelling Bluetooth headphones with 30-hour battery life."}, + {"name": "Mechanical Keyboard", "category": "Electronics", "price": 99.99, "description": "Compact TKL mechanical keyboard with tactile switches and RGB backlight."}, + {"name": "USB-C Hub", "category": "Electronics", "price": 39.99, "description": "7-in-1 USB-C hub with HDMI, SD card, and 100W power delivery."}, + {"name": "Webcam 1080p", "category": "Electronics", "price": 59.99, "description": "Full HD webcam with built-in microphone for video conferencing."}, + {"name": "Portable SSD 1TB", "category": "Electronics", "price": 109.99, "description": "Ultra-fast portable SSD with USB 3.2 Gen 2 and rugged aluminum case."}, + {"name": "Smart Watch", "category": "Electronics", "price": 199.99, "description": "Fitness smart watch with heart-rate monitor and 7-day battery."}, + {"name": "Bluetooth Speaker", "category": "Electronics", "price": 49.99, "description": "Waterproof portable speaker with 360° sound and 12-hour battery."}, + {"name": "Laptop Stand", "category": "Electronics", "price": 29.99, "description": "Adjustable aluminum laptop stand compatible with 10-16 inch laptops."}, + # Clothing + {"name": "Running Shoes", "category": "Clothing", "price": 89.99, "description": "Lightweight breathable running shoes with cushioned sole and wide toe box."}, + {"name": "Merino Wool Sweater", "category": "Clothing", "price": 69.99, "description": "Soft and warm 100% merino wool crewneck sweater, machine washable."}, + {"name": "Waterproof Jacket", "category": "Clothing", "price": 129.99, "description": "Lightweight packable rain jacket with taped seams and adjustable hood."}, + {"name": "Slim-Fit Chinos", "category": "Clothing", "price": 54.99, "description": "Stretch slim-fit chino trousers available in multiple neutral colors."}, + {"name": "Cotton T-Shirt 3-Pack", "category": "Clothing", "price": 29.99, "description": "Everyday crewneck cotton tees in classic white, black, and grey."}, + {"name": "Leather Belt", "category": "Clothing", "price": 34.99, "description": "Full-grain leather dress belt with silver-tone pin buckle."}, + {"name": "Hiking Boots", "category": "Clothing", "price": 149.99, "description": "Waterproof mid-cut hiking boots with vibram sole and ankle support."}, + {"name": "Puffer Vest", "category": "Clothing", "price": 59.99, "description": "Lightweight down-fill vest with snap front closure and two hand pockets."}, + # Books + {"name": "Deep Learning Book", "category": "Books", "price": 44.99, "description": "Comprehensive guide to deep learning foundations and modern architectures."}, + {"name": "Python Cookbook", "category": "Books", "price": 39.99, "description": "Practical recipes for writing idiomatic, modern Python code."}, + {"name": "Distributed Systems", "category": "Books", "price": 49.99, "description": "Principles and patterns of building reliable distributed systems at scale."}, + {"name": "Clean Code", "category": "Books", "price": 34.99, "description": "A handbook of agile software craftsmanship principles and best practices."}, + {"name": "The Algorithm Design", "category": "Books", "price": 59.99, "description": "Classic reference covering algorithm design techniques with worked examples."}, + {"name": "Designing Data-Intensive Apps", "category": "Books", "price": 54.99, "description": "In-depth look at the systems behind modern data-intensive applications."}, + # Home & Garden + {"name": "French Press Coffee", "category": "Home & Garden", "price": 29.99, "description": "Stainless steel French press with double-wall insulation, 34 oz."}, + {"name": "Cast Iron Skillet", "category": "Home & Garden", "price": 39.99, "description": "Pre-seasoned 10-inch cast iron skillet suitable for all cooktops."}, + {"name": "Bamboo Cutting Board", "category": "Home & Garden", "price": 24.99, "description": "Large reversible bamboo cutting board with juice groove and handle."}, + {"name": "Air Purifier", "category": "Home & Garden", "price": 119.99, "description": "True HEPA air purifier for rooms up to 500 sq ft, whisper-quiet."}, + {"name": "Succulent Set 6-Pack", "category": "Home & Garden", "price": 19.99, "description": "Assorted live succulents in 2-inch pots, easy-care indoor plants."}, + {"name": "Scented Soy Candle", "category": "Home & Garden", "price": 18.99, "description": "Hand-poured lavender and vanilla soy wax candle, 50-hour burn time."}, + # Sports + {"name": "Yoga Mat", "category": "Sports", "price": 34.99, "description": "6mm thick non-slip yoga mat with carry strap, eco-friendly TPE."}, + {"name": "Resistance Bands Set", "category": "Sports", "price": 24.99, "description": "Set of 5 fabric resistance bands with varying tension levels."}, + {"name": "Jump Rope", "category": "Sports", "price": 14.99, "description": "Adjustable speed jump rope with ball-bearing handles and steel cable."}, + {"name": "Water Bottle 32oz", "category": "Sports", "price": 29.99, "description": "Insulated stainless steel water bottle keeps drinks cold 24 hr / hot 12 hr."}, + {"name": "Foam Roller", "category": "Sports", "price": 22.99, "description": "High-density foam roller for muscle recovery and myofascial release."}, + {"name": "Dumbbell Set", "category": "Sports", "price": 89.99, "description": "Adjustable dumbbell set 5-25 lb per hand with compact storage tray."}, +] + +IMAGE_SIZE = (224, 224) + +# Adjectives and category-specific suffixes used to generate product variants +_VARIANT_ADJECTIVES = [ + "Premium", "Pro", "Lite", "Ultra", "Essential", "Plus", + "Max", "Mini", "Elite", "Classic", "Advanced", "Standard", + "Signature", "Sport", "Studio", "Travel", +] + +_VARIANT_SUFFIXES = { + "Electronics": [ + "Includes extended warranty.", + "Optimized for remote work.", + "Designed for content creators.", + "Features USB-C connectivity.", + "With advanced noise reduction.", + "Smart power management included.", + ], + "Clothing": [ + "Available in 8 colors.", + "Made from sustainable materials.", + "Reinforced stitching throughout.", + "Features moisture-wicking fabric.", + "UPF 50+ sun protection.", + "Slim-fit design.", + ], + "Books": [ + "Includes digital access code.", + "Updated 2024 edition.", + "With practice exercises throughout.", + "Annotated edition.", + "Includes 100+ code examples.", + "With online companion resources.", + ], + "Home & Garden": [ + "Dishwasher safe.", + "BPA-free construction.", + "Lifetime warranty included.", + "Eco-friendly packaging.", + "NSF certified.", + "Made from recycled materials.", + ], + "Sports": [ + "With antimicrobial coating.", + "Suitable for all fitness levels.", + "Includes carrying bag.", + "Made from recycled materials.", + "Ergonomic grip design.", + "Machine washable.", + ], +} + + +# --------------------------------------------------------------------------- +# Synthetic image generation +# --------------------------------------------------------------------------- + +def make_product_image(product: Dict, seed: int = 0) -> np.ndarray: + """Generate a synthetic product image as a numpy array (H, W, 3) uint8. + + The image is a solid-colored rectangle with the product name overlaid, + so it's visually distinct per category without requiring real photos. + """ + rng = random.Random(seed) + base_color = CATEGORY_COLORS[product["category"]] + + # Add slight per-product color jitter + color = tuple( + max(0, min(255, c + rng.randint(-30, 30))) for c in base_color + ) + + img = Image.new("RGB", IMAGE_SIZE, color=color) + draw = ImageDraw.Draw(img) + + # Draw a centered rounded rectangle as a "product card" + margin = 20 + draw.rounded_rectangle( + [margin, margin, IMAGE_SIZE[0] - margin, IMAGE_SIZE[1] - margin], + radius=15, + fill=tuple(max(0, c - 40) for c in color), + ) + + # Write product name (word-wrap at ~15 chars) + words = product["name"].split() + lines, line = [], [] + for word in words: + if sum(len(w) for w in line) + len(line) + len(word) <= 15: + line.append(word) + else: + lines.append(" ".join(line)) + line = [word] + if line: + lines.append(" ".join(line)) + + # Try to use a basic font; fall back to default + try: + font = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", 20) + small_font = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", 13) + except Exception: + font = ImageFont.load_default() + small_font = font + + total_h = len(lines) * 24 + y = (IMAGE_SIZE[1] - total_h) // 2 - 10 + for line_text in lines: + bbox = draw.textbbox((0, 0), line_text, font=font) + w = bbox[2] - bbox[0] + draw.text(((IMAGE_SIZE[0] - w) // 2, y), line_text, fill="white", font=font) + y += 24 + + # Category label at bottom + cat = product["category"] + bbox = draw.textbbox((0, 0), cat, font=small_font) + w = bbox[2] - bbox[0] + draw.text(((IMAGE_SIZE[0] - w) // 2, IMAGE_SIZE[1] - 35), cat, fill="white", font=small_font) + + return np.array(img) + + +_DEMO_IMAGES_DIR = Path(__file__).parent.parent / "data" / "demo_images" + + +def get_product_image(product: Dict) -> np.ndarray: + """Load the bundled realistic product image, falling back to synthetic. + + Variant products created by ``expand_catalog`` carry a ``_base_name`` key + pointing to their base product so they reuse that product's real image. + """ + image_name = product.get("_base_name", product["name"]) + safe_name = image_name.replace(" ", "_").replace("/", "-") + path = _DEMO_IMAGES_DIR / f"{safe_name}.jpg" + if path.exists(): + return np.array(Image.open(path).convert("RGB")) + return make_product_image(product, seed=hash(product["name"]) % 1000) + + +def image_to_bytes(img_array: np.ndarray, fmt: str = "JPEG") -> bytes: + """Convert numpy (H,W,3) uint8 array to compressed image bytes.""" + buf = io.BytesIO() + Image.fromarray(img_array).save(buf, format=fmt) + return buf.getvalue() + + +def bytes_to_image(raw: bytes) -> np.ndarray: + """Decode compressed image bytes back to numpy (H,W,3) uint8.""" + return np.array(Image.open(io.BytesIO(raw)).convert("RGB")) + + +# --------------------------------------------------------------------------- +# Text utilities +# --------------------------------------------------------------------------- + +def clean_text(text: str) -> str: + """Basic text cleaning: strip, lowercase, collapse whitespace.""" + import re + text = text.strip().lower() + text = re.sub(r"\s+", " ", text) + return text + + +def make_training_text(product: Dict) -> str: + """Combine product fields into a single string for embedding training.""" + return f"{product['name']}. {product['description']}" + + +# --------------------------------------------------------------------------- +# Dataset generation +# --------------------------------------------------------------------------- + +def generate_catalog( + products: Optional[List[Dict]] = None, + output_dir: str = "data/raw", + seed: int = 42, +) -> List[Dict]: + """Generate the synthetic product catalog and save images to disk. + + Returns a list of records ready to be loaded into a Ray Dataset. + """ + if products is None: + products = PRODUCTS + + Path(output_dir).mkdir(parents=True, exist_ok=True) + + records = [] + for i, p in enumerate(products): + img_array = get_product_image(p) + img_bytes = image_to_bytes(img_array) + + record = { + "product_id": f"P{i:04d}", + "name": p["name"], + "category": p["category"], + "description": p["description"], + "price": p["price"], + "image_bytes": img_bytes, + "training_text": make_training_text(p), + } + records.append(record) + + print(f"Generated {len(records)} products in '{output_dir}'") + return records + + +def expand_catalog( + base_products: Optional[List[Dict]] = None, + target_size: int = 1000, + seed: int = 42, +) -> List[Dict]: + """Expand the base catalog to *target_size* by generating product variants. + + Variant products share images with their base product via ``_base_name``, + but have distinct names and descriptions — giving the embedding model + diverse training signal without requiring additional real photos. + """ + if base_products is None: + base_products = PRODUCTS + if not base_products: + raise ValueError("base_products cannot be empty") + + rng = random.Random(seed) + expanded = [dict(p) for p in base_products] # start with originals + + i = 0 + while len(expanded) < target_size: + base = base_products[i % len(base_products)] + adj = rng.choice(_VARIANT_ADJECTIVES) + suffix = rng.choice(_VARIANT_SUFFIXES[base["category"]]) + expanded.append({ + "name": f"{adj} {base['name']}", + "category": base["category"], + "price": round(base["price"] * rng.uniform(0.75, 1.4), 2), + "description": base["description"].rstrip(".") + ". " + suffix, + "_base_name": base["name"], + }) + i += 1 + + return expanded[:target_size] + + +# --------------------------------------------------------------------------- +# Preprocessing helpers (used inside Ray Data map_batches) +# --------------------------------------------------------------------------- + +def preprocess_image_batch(batch: Dict) -> Dict: + """Normalize image bytes to float32 tensor bytes (224x224x3, range [0,1]). + + Suitable as a Ray Data map_batches function. + """ + processed = [] + for raw in batch["image_bytes"]: + img = np.array( + Image.open(io.BytesIO(raw)).convert("RGB").resize(IMAGE_SIZE) + ).astype(np.float32) / 255.0 + # Store as bytes (float32 little-endian) to keep Parquet schema simple + processed.append(img.tobytes()) + batch["image_tensor_bytes"] = processed + return batch + + +def preprocess_text_batch(batch: Dict) -> Dict: + """Clean training text. Tokenization happens inside the trainer.""" + batch["text_clean"] = [clean_text(t) for t in batch["training_text"]] + return batch + + +def decode_image_tensor(raw: bytes) -> np.ndarray: + """Inverse of preprocess_image_batch: bytes -> float32 (224,224,3).""" + return np.frombuffer(raw, dtype=np.float32).reshape(*IMAGE_SIZE, 3) + + +# --------------------------------------------------------------------------- +# Ray helpers +# --------------------------------------------------------------------------- + +def init_ray() -> None: + """Initialize Ray with reduced logging (idempotent).""" + import logging + import ray + import ray.data + + for name in ["ray", "ray.data", "ray.train", "ray.tune", "ray.serve", + "ray._private", "ray.runtime_env"]: + logging.getLogger(name).setLevel(logging.WARNING) + + ray.init( + ignore_reinit_error=True, + log_to_driver=False, + logging_level=logging.WARNING, + ) + ray.data.DataContext.get_current().enable_progress_bars = True + + +# --------------------------------------------------------------------------- +# Notebook convenience helpers +# --------------------------------------------------------------------------- + +def attach_clean_text(records: List[Dict]) -> List[Dict]: + """Add a `text_clean` field to each record (in-place), derived from + `training_text`. Same cleaning rules as :func:`clean_text`. + """ + for r in records: + r["text_clean"] = clean_text(r["training_text"]) + return records + + +def sample_per_category( + records: List[Dict], + n_per_category: int, + seed: int = 42, +) -> List[Dict]: + """Return a class-balanced sample — up to *n_per_category* items per + `category` — which gives contrastive training enough positives for every + category even when the dataset is large. + """ + from collections import defaultdict + + buckets: Dict[str, List[Dict]] = defaultdict(list) + for r in records: + buckets[r["category"]].append(r) + + rng = random.Random(seed) + sampled: List[Dict] = [] + for bucket in buckets.values(): + k = min(n_per_category, len(bucket)) + sampled.extend(rng.sample(bucket, k)) + return sampled + + +def resolve_artifact_paths(here: Optional[str] = None) -> Dict[str, str]: + """Pick where to read/write models + embeddings. + + On Anyscale clusters we prefer `/mnt/cluster_storage` so every node sees + the same files; on a laptop we fall back to a local `models/` folder. + Returns a dict with keys: model_dir, embeddings_path, metadata_path, + train_result_dir. + """ + here = here or os.path.abspath(".") + shared = "/mnt/cluster_storage" + use_shared = os.path.isdir(shared) + base = shared if use_shared else os.path.join(here, "models") + + def _p(shared_name: str, local_name: str) -> str: + return ( + os.path.join(shared, shared_name) + if use_shared + else os.path.join(here, "models", local_name) + ) + + return { + "model_dir": _p("ecomm_embedding_model", "embedding_model"), + "embeddings_path": _p("ecomm_product_embeddings.npy", "product_embeddings.npy"), + "metadata_path": _p("ecomm_product_metadata.json", "product_metadata.json"), + "train_result_dir": _p("ecomm_ray_train_results", "ray_train_results"), + } + + +def encode_image_base64(image_array: np.ndarray, fmt: str = "JPEG") -> str: + """Encode a (H,W,3) numpy image to a base64 ASCII string — the payload + shape our `/recommend` endpoint expects. + """ + import base64 + return base64.b64encode(image_to_bytes(image_array, fmt=fmt)).decode() + + +def post_recommend( + image_array: np.ndarray, + url: str = "http://localhost:8000/recommend", +) -> Dict: + """POST an image to the recommendation endpoint and return the JSON body.""" + import requests + + payload = {"image_base64": encode_image_base64(image_array)} + return requests.post(url, json=payload).json() diff --git a/templates/ecommerce_end_to_end/utils/embedding.py b/templates/ecommerce_end_to_end/utils/embedding.py new file mode 100644 index 000000000..6e9ca947f --- /dev/null +++ b/templates/ecommerce_end_to_end/utils/embedding.py @@ -0,0 +1,126 @@ +"""Batch embedding utilities for the product catalog.""" + +import json +import time +from pathlib import Path +from typing import List, Dict + +import numpy as np +import ray + + +def embed_catalog_in_parallel( + ds: "ray.data.Dataset", + model_dir: str, + pool_size: int = 2, + batch_size: int = 8, +) -> List[Dict]: + """Run the fine-tuned model across `ds` with an actor pool. + + Each actor loads the model exactly once and then streams batches through + it, so cost amortises across the whole catalog instead of per-row. + Scaling up is a matter of bumping `pool_size`. + """ + assert Path(model_dir).exists() and any(Path(model_dir).iterdir()), ( + f"Model not found at {model_dir} — run the fine-tuning stage first" + ) + return ds.map_batches( + ProductEmbedder, + fn_constructor_kwargs={"model_dir": model_dir}, + batch_size=batch_size, + num_cpus=1, + compute=ray.data.ActorPoolStrategy(size=pool_size), + batch_format="numpy", + ).take_all() + + +def save_embeddings_and_metadata( + rows: List[Dict], + embeddings_path: str, + metadata_path: str, +) -> tuple: + """Split the embedder output into a dense vector matrix (saved as .npy) + and a small JSON sidecar with just the fields the serving layer needs. + Keeping them separate avoids loading product text into memory at query time. + """ + embeddings = np.array([r["embedding"] for r in rows]) + metadata = [ + {"product_id": r["product_id"], "name": r["name"], "category": r["category"]} + for r in rows + ] + Path(embeddings_path).parent.mkdir(parents=True, exist_ok=True) + np.save(embeddings_path, embeddings) + with open(metadata_path, "w") as f: + json.dump(metadata, f, indent=2) + return embeddings, metadata + + +class ProductEmbedder: + """Ray Data actor: loads the fine-tuned model once, encodes batches of text.""" + + def __init__(self, model_dir: str): + from sentence_transformers import SentenceTransformer + self.model = SentenceTransformer(model_dir) + + def __call__(self, batch: dict) -> dict: + texts = batch["text_clean"].tolist() + embs = self.model.encode(texts, convert_to_numpy=True, show_progress_bar=False) + return { + "product_id": batch["product_id"], + "name": batch["name"], + "category": batch["category"], + "embedding": list(embs), + } + + +def run_batch_embedding( + preprocessed_dir: str, + model_dir: str, + embeddings_path: str, + metadata_path: str, + batch_size: int = 8, +) -> tuple: + """Run Stage 3: embed entire catalog with fine-tuned model. + + Returns: + (embeddings, metadata) — numpy array of shape (N, D) and list of dicts. + """ + t0 = time.time() + print("=" * 60) + print("STAGE 3 — RAY DATA: BATCH EMBEDDING") + print("=" * 60) + + ds = ( + ray.data.read_parquet(preprocessed_dir) + .select_columns(["product_id", "name", "category", "text_clean"]) + ) + print(f"\nRows to embed: {ds.count()}") + print("Running batch embedding …") + + rows = ( + ds.map_batches( + ProductEmbedder, + fn_constructor_kwargs={"model_dir": model_dir}, + batch_size=batch_size, + num_cpus=1, + compute=ray.data.ActorPoolStrategy(size=2), + batch_format="numpy", + ) + .take_all() + ) + + embeddings = np.array([r["embedding"] for r in rows]) + metadata = [ + {"product_id": r["product_id"], "name": r["name"], "category": r["category"]} + for r in rows + ] + + np.save(embeddings_path, embeddings) + with open(metadata_path, "w") as f: + json.dump(metadata, f, indent=2) + + print(f"\nBatch embedding complete! ({time.time()-t0:.1f}s)") + print(f" Embeddings : {embeddings_path} shape={embeddings.shape}") + print(f" Metadata : {metadata_path} ({len(metadata)} products)") + print("=" * 60) + return embeddings, metadata diff --git a/templates/ecommerce_end_to_end/utils/stages.py b/templates/ecommerce_end_to_end/utils/stages.py new file mode 100644 index 000000000..bc78998ce --- /dev/null +++ b/templates/ecommerce_end_to_end/utils/stages.py @@ -0,0 +1,142 @@ +"""High-level stage runner functions used by the demo notebook.""" + +import os +import shutil +import time +from pathlib import Path + +import ray + + +# --------------------------------------------------------------------------- +# Stage 1 — Ray Data: Preprocessing +# --------------------------------------------------------------------------- + +def run_preprocessing(output_dir: str, batch_size: int = 8) -> int: + """Generate catalog, preprocess images + text, write Parquet shards. + + Returns the number of rows written. + """ + from . import PRODUCTS, generate_catalog, preprocess_image_batch, preprocess_text_batch + + t0 = time.time() + print("=" * 60) + print("STAGE 1 — RAY DATA: PREPROCESSING") + print("=" * 60) + + raw_dir = str(Path(output_dir).parent / "raw") + + print("\n[1/4] Generating product catalog …") + records = generate_catalog(products=PRODUCTS, output_dir=raw_dir) + ds = ray.data.from_items(records) + print(f" Rows: {ds.count()}") + + print("\n[2/4] Preprocessing images (resize + normalise) …") + ds = ds.map_batches( + preprocess_image_batch, batch_size=batch_size, num_cpus=1, batch_format="numpy" + ) + + print("\n[3/4] Preprocessing text (clean) …") + ds = ds.map_batches( + preprocess_text_batch, batch_size=batch_size, num_cpus=1, batch_format="numpy" + ) + + print(f"\n[4/4] Writing to '{output_dir}' …") + if os.path.exists(output_dir): + shutil.rmtree(output_dir) + Path(output_dir).mkdir(parents=True, exist_ok=True) + ds.write_parquet(output_dir) + + n = ray.data.read_parquet(output_dir).count() + print(f"\nPreprocessing complete! Rows: {n} ({time.time()-t0:.1f}s)") + print("=" * 60) + return n + + +# --------------------------------------------------------------------------- +# Stage 2 — Ray Train: Embedding Fine-Tuning +# --------------------------------------------------------------------------- + +def run_training( + preprocessed_dir: str, + model_output_dir: str, + train_result_dir: str, + base_model: str = "sentence-transformers/all-MiniLM-L6-v2", + num_epochs: int = 2, + batch_size: int = 8, + lr: float = 2e-5, + seed: int = 42, +): + """Fine-tune the embedding model with contrastive loss via TorchTrainer. + + Saves the best checkpoint to *model_output_dir* and returns the Ray Train + ``Result`` object (includes ``metrics_dataframe`` for plotting). + """ + import glob + + import torch + from ray.train import CheckpointConfig, FailureConfig, RunConfig, ScalingConfig + from ray.train.torch import TorchTrainer + from sentence_transformers import SentenceTransformer + + from .training import train_loop_per_worker + + parquet_files = glob.glob(os.path.join(preprocessed_dir, "*.parquet")) + if not parquet_files: + raise FileNotFoundError( + f"No parquet files found in '{preprocessed_dir}'. " + "Run Stage 1 (preprocessing) first." + ) + + t0 = time.time() + print("=" * 60) + print("STAGE 2 — RAY TRAIN: EMBEDDING FINE-TUNING") + print("=" * 60) + + records = ( + ray.data.read_parquet(preprocessed_dir) + .select_columns(["product_id", "name", "category", "text_clean"]) + .take_all() + ) + print(f"\nLoaded {len(records)} records for training") + + Path(train_result_dir).mkdir(parents=True, exist_ok=True) + + trainer = TorchTrainer( + train_loop_per_worker=train_loop_per_worker, + train_loop_config={ + "base_model": base_model, + "epochs": num_epochs, + "batch_size": batch_size, + "lr": lr, + "seed": seed, + "records": records, + }, + scaling_config=ScalingConfig( + num_workers=1, + use_gpu=torch.cuda.is_available(), + ), + run_config=RunConfig( + name="ecomm_embedding_finetune", + storage_path=os.path.abspath(train_result_dir), + checkpoint_config=CheckpointConfig( + num_to_keep=2, + checkpoint_score_attribute="train_loss", + checkpoint_score_order="min", + ), + failure_config=FailureConfig(max_failures=1), + ), + ) + + result = trainer.fit() + print(f"\nTraining complete! ({time.time()-t0:.1f}s)") + + print("\nSaving fine-tuned model …") + Path(model_output_dir).mkdir(parents=True, exist_ok=True) + checkpoint = result.best_checkpoints[0][0] + with checkpoint.as_directory() as ckpt_dir: + model = SentenceTransformer(ckpt_dir) + model.save(model_output_dir) + print(f" Model saved : {model_output_dir}") + print("=" * 60) + return result diff --git a/templates/ecommerce_end_to_end/utils/training.py b/templates/ecommerce_end_to_end/utils/training.py new file mode 100644 index 000000000..915f4ad32 --- /dev/null +++ b/templates/ecommerce_end_to_end/utils/training.py @@ -0,0 +1,196 @@ +"""Training utilities for embedding fine-tuning with contrastive loss.""" + +import os +import tempfile +from pathlib import Path +from typing import List, Dict + +import numpy as np +import torch +import torch.nn.functional as F +from torch.utils.data import DataLoader, Dataset +from ray.train import Checkpoint, get_checkpoint, get_context, report + +SEED = 42 + + +def build_embedding_trainer( + records: List[Dict], + train_result_dir: str, + base_model: str = "sentence-transformers/all-MiniLM-L6-v2", + epochs: int = 2, + batch_size: int = 8, + lr: float = 2e-5, + seed: int = SEED, +): + """Build a Ray :class:`TorchTrainer` configured for contrastive fine-tuning. + + Keeps checkpoint/failure config next to the training loop so the notebook + only has to say *what* it wants to train, not *how* Ray Train wires it up. + """ + from ray.train import ( + CheckpointConfig, + FailureConfig, + RunConfig, + ScalingConfig, + ) + from ray.train.torch import TorchTrainer + + return TorchTrainer( + train_loop_per_worker=train_loop_per_worker, + train_loop_config={ + "base_model": base_model, + "epochs": epochs, + "batch_size": batch_size, + "lr": lr, + "seed": seed, + "records": records, + }, + scaling_config=ScalingConfig( + num_workers=1, + use_gpu=torch.cuda.is_available(), + ), + run_config=RunConfig( + name="ecomm_embedding_finetune", + storage_path=os.path.abspath(train_result_dir), + checkpoint_config=CheckpointConfig( + num_to_keep=2, + checkpoint_score_attribute="train_loss", + checkpoint_score_order="min", + ), + failure_config=FailureConfig(max_failures=1), + ), + ) + + +def save_best_sentence_transformer(result, output_dir: str) -> str: + """Restore the lowest-loss checkpoint from a Ray Train `Result` and save it + in SentenceTransformer format — i.e. a folder that `SentenceTransformer(path)` + can load directly for serving. + """ + from sentence_transformers import SentenceTransformer + + Path(output_dir).mkdir(parents=True, exist_ok=True) + best_checkpoint = result.best_checkpoints[0][0] + with best_checkpoint.as_directory() as ckpt_dir: + SentenceTransformer(ckpt_dir).save(output_dir) + return output_dir + + +class ContrastivePairDataset(Dataset): + """ + (anchor, positive, 1.0) for same-category pairs; + (anchor, negative, -1.0) for cross-category pairs. + Labels match CosineSimilarityLoss / MSE-on-cosine expectations. + """ + + def __init__(self, records, neg_ratio=0.5, seed=SEED): + rng = np.random.default_rng(seed) + cat_to_idx = {} + for i, r in enumerate(records): + cat_to_idx.setdefault(r["category"], []).append(i) + + pairs = [] + for indices in cat_to_idx.values(): + for i, a in enumerate(indices): + for b in indices[i + 1 :]: + pairs.append( + (records[a]["text_clean"], records[b]["text_clean"], 1.0) + ) + + cats = list(cat_to_idx.keys()) + n_neg = max(1, int(len(pairs) * neg_ratio)) + for _ in range(n_neg): + cat_a, cat_b = rng.choice(cats, size=2, replace=False) + ia = rng.choice(cat_to_idx[cat_a]) + ib = rng.choice(cat_to_idx[cat_b]) + pairs.append((records[ia]["text_clean"], records[ib]["text_clean"], -1.0)) + + rng.shuffle(pairs) + self.pairs = pairs + + def __len__(self): + return len(self.pairs) + + def __getitem__(self, idx): + a, b, label = self.pairs[idx] + return a, b, torch.tensor(label, dtype=torch.float32) + + +def _forward_embeddings(model, texts, device): + """Run a SentenceTransformer forward pass with gradient tracking.""" + features = model.tokenize(texts) + features = { + k: v.to(device) for k, v in features.items() if isinstance(v, torch.Tensor) + } + return model(features)["sentence_embedding"] + + +def train_loop_per_worker(config: dict) -> None: + """Fine-tune the embedding model on each distributed Ray Train worker.""" + from sentence_transformers import SentenceTransformer + + rank = get_context().get_world_rank() + device = "cuda" if torch.cuda.is_available() else "cpu" + records = config["records"] + + pair_dataset = ContrastivePairDataset(records, seed=config["seed"]) + + def collate(batch): + texts_a, texts_b, labels = zip(*batch) + return list(texts_a), list(texts_b), torch.stack(labels) + + loader = DataLoader( + pair_dataset, + batch_size=config["batch_size"], + shuffle=True, + collate_fn=collate, + ) + + model = SentenceTransformer(config["base_model"], device=device) + optimizer = torch.optim.AdamW( + model.parameters(), lr=config["lr"], weight_decay=0.01 + ) + + start_epoch = 0 + ckpt = get_checkpoint() + if ckpt: + with ckpt.as_directory() as d: + meta = torch.load( + os.path.join(d, "meta.pt"), map_location="cpu", weights_only=False + ) + start_epoch = meta.get("epoch", 0) + 1 + model = SentenceTransformer(d, device=device) + + if rank == 0: + print( + f"Fine-tuning {config['base_model']} on {device} " + f"({config['epochs']} epochs, {len(pair_dataset)} pairs)" + ) + + for epoch in range(start_epoch, config["epochs"]): + model.train() + total_loss, n_batches = 0.0, 0 + for texts_a, texts_b, labels in loader: + labels = labels.to(device) + emb_a = _forward_embeddings(model, texts_a, device) + emb_b = _forward_embeddings(model, texts_b, device) + loss = F.mse_loss(F.cosine_similarity(emb_a, emb_b), labels) + optimizer.zero_grad() + loss.backward() + optimizer.step() + total_loss += loss.item() + n_batches += 1 + + avg_loss = total_loss / max(1, n_batches) + if rank == 0: + print(f" Epoch {epoch+1:2d}/{config['epochs']} loss={avg_loss:.4f}") + + with tempfile.TemporaryDirectory() as tmpdir: + if rank == 0: + model.save(tmpdir) + torch.save({"epoch": epoch}, os.path.join(tmpdir, "meta.pt")) + ckpt_out = Checkpoint.from_directory(tmpdir) + else: + ckpt_out = None + report({"epoch": epoch, "train_loss": avg_loss}, checkpoint=ckpt_out) diff --git a/templates/ecommerce_end_to_end/utils/viz.py b/templates/ecommerce_end_to_end/utils/viz.py new file mode 100644 index 000000000..8fcf00fc3 --- /dev/null +++ b/templates/ecommerce_end_to_end/utils/viz.py @@ -0,0 +1,288 @@ +"""Visualization utilities for embedding analysis.""" + +import os +from collections import Counter +from typing import Callable, Dict, List, Optional, Tuple, Union + +import numpy as np +import matplotlib.pyplot as plt + + +def plot_category_samples(products: List[Dict], categories: List[str], get_image_fn: Callable) -> None: + """Show one sample product image per category.""" + fig, axes = plt.subplots(1, len(categories), figsize=(15, 3)) + for ax, cat in zip(axes, categories): + product = next(p for p in products if p["category"] == cat) + img = get_image_fn(product) + ax.imshow(img) + ax.set_title(cat, fontsize=9) + ax.axis("off") + plt.suptitle("One sample image per category", fontsize=12) + plt.tight_layout() + plt.show() + + +def plot_products_per_category(products: List[Dict]) -> None: + """Horizontal bar chart of product counts per category.""" + cats = [p["category"] for p in products] + counts = sorted(Counter(cats).items()) + plt.barh(*zip(*counts)) + plt.xlabel("Products") + plt.title("Products per category") + plt.tight_layout() + plt.show() + + +def plot_training_loss(metrics_dataframe) -> None: + """Line plot of training loss across epochs.""" + m = metrics_dataframe + plt.plot(m["epoch"] + 1, m["train_loss"], marker="o") + plt.xlabel("Epoch") + plt.ylabel("Loss") + plt.title("Training Loss") + plt.tight_layout() + plt.show() + + +def plot_recommendations( + query_img: np.ndarray, + recommendations: List[Dict], + products: List[Dict], + caption: str, + get_image_fn: Callable, +) -> None: + """Show query image alongside recommended products.""" + fig, axes = plt.subplots(1, len(recommendations) + 1, figsize=(18, 3)) + + axes[0].imshow(query_img) + axes[0].set_title("Query image", fontsize=9, color="navy", fontweight="bold") + axes[0].axis("off") + + all_prods = {p["name"]: p for p in products} + for ax, rec in zip(axes[1:], recommendations): + prod_name = rec["name"] + prod = all_prods.get(prod_name, {"name": prod_name, "category": rec["category"]}) + rec_img = get_image_fn({**prod, "category": rec["category"]}) + ax.imshow(rec_img) + ax.set_title( + f"{rec['rank']}. {rec['name'][:18]}\n{rec['similarity']:.2f}", fontsize=8 + ) + ax.axis("off") + + plt.suptitle(f'Caption: "{caption}"', fontsize=11) + plt.tight_layout() + plt.show() + + +def plot_similarity_heatmap( + embeddings, labels, n=10, title="Product embedding similarity (top 10)" +): + """Plot cosine similarity heatmap for the first n embeddings.""" + n = min(n, len(embeddings)) + sub = embeddings[:n] + norms = np.linalg.norm(sub, axis=1, keepdims=True) + sim = (sub / norms) @ (sub / norms).T + + short_labels = [lbl[:20] for lbl in labels[:n]] + fig, ax = plt.subplots(figsize=(8, 6)) + im = ax.imshow(sim, vmin=-1, vmax=1, cmap="RdYlGn") + ax.set_xticks(range(n)) + ax.set_xticklabels(short_labels, rotation=45, ha="right", fontsize=8) + ax.set_yticks(range(n)) + ax.set_yticklabels(short_labels, fontsize=8) + plt.colorbar(im, ax=ax, label="Cosine similarity") + ax.set_title(title) + plt.tight_layout() + plt.show() + + +def _plot_tsne_side_by_side( + base_embs: np.ndarray, + tuned_embs: np.ndarray, + labels: List[str], + *, + base_title: str = "Base model (all-MiniLM-L6-v2)", + tuned_title: str = "Fine-tuned (contrastive loss)", +) -> None: + """Side-by-side t-SNE of two embedding matrices, coloured by category labels.""" + from sklearn.manifold import TSNE + + cat_list = sorted(set(labels)) + colors = [cat_list.index(c) for c in labels] + cmap = plt.cm.get_cmap("tab10", len(cat_list)) + + fig, axes = plt.subplots(1, 2, figsize=(14, 5)) + for ax, embs, title in [ + (axes[0], base_embs, base_title), + (axes[1], tuned_embs, tuned_title), + ]: + xy = TSNE(n_components=2, perplexity=10, random_state=42).fit_transform(embs) + ax.scatter( + xy[:, 0], xy[:, 1], c=colors, cmap=cmap, s=80, vmin=0, vmax=len(cat_list) + ) + ax.set_title(title) + ax.axis("off") + + handles = [ + plt.Line2D( + [0], + [0], + marker="o", + color="w", + markerfacecolor=cmap(i), + markersize=9, + label=c, + ) + for i, c in enumerate(cat_list) + ] + fig.legend(handles=handles, loc="lower center", ncol=len(cat_list), fontsize=9) + plt.suptitle("t-SNE: product embeddings coloured by category", fontsize=12) + plt.tight_layout(rect=[0, 0.06, 1, 1]) + plt.show() + + +def plot_tsne_comparison( + metadata_or_base: Union[List[Dict], np.ndarray], + embeddings_or_tuned: np.ndarray, + labels: Optional[List[str]] = None, + *, + catalog_records: Optional[List[Dict]] = None, + base_model_name: str = "sentence-transformers/all-MiniLM-L6-v2", +) -> Optional[Tuple[np.ndarray, np.ndarray, List[Dict]]]: + """Compare base MiniLM vs fine-tuned embeddings with a t-SNE plot. + + **High-level (recommended):** pass ``metadata``, ``tuned_embeddings``, and + ``catalog_records`` (list of dicts with ``product_id`` and ``text_clean``). + Aligns rows, encodes text with the base model, plots, and returns + ``(base_embs, tuned_embs, aligned_metadata)`` for downstream metrics. + + **Low-level:** pass ``base_embs``, ``tuned_embs``, and ``labels`` (three + positional args) to plot precomputed arrays only; returns ``None``. + """ + if labels is not None: + _plot_tsne_side_by_side( + np.asarray(metadata_or_base), + np.asarray(embeddings_or_tuned), + labels, + ) + return None + + from sentence_transformers import SentenceTransformer + + metadata = metadata_or_base # type: ignore[assignment] + tuned_embeddings = np.asarray(embeddings_or_tuned) + + if catalog_records is None: + raise ValueError("catalog_records is required") + id_to_text = {r["product_id"]: r["text_clean"] for r in catalog_records} + + idx = [i for i, m in enumerate(metadata) if m["product_id"] in id_to_text] + if not idx: + raise ValueError("No metadata rows overlap catalog_records product_id values.") + + aligned_metadata = [metadata[i] for i in idx] + ft_embs = tuned_embeddings[idx] + texts = [id_to_text[m["product_id"]] for m in aligned_metadata] + + base_embs = SentenceTransformer(base_model_name).encode( + texts, convert_to_numpy=True + ) + category_labels = [m["category"] for m in aligned_metadata] + _plot_tsne_side_by_side(base_embs, ft_embs, category_labels) + return base_embs, ft_embs, aligned_metadata + + +def compute_category_precision_at_k(embeddings, metadata, k: int = 5): + """Share of each product's k nearest neighbors (by cosine similarity) in the same category. + + Often hits 100% when product text already matches its category; use + ``mean_rank_first_cross_category_neighbor`` or ``mean_intra_inter_sim_gap`` + to compare base vs fine-tuned models in that case. + """ + n = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True) + sim = n @ n.T + np.fill_diagonal(sim, -np.inf) + cats = [m["category"] for m in metadata] + return np.mean( + [ + sum(cats[j] == cats[i] for j in np.argsort(sim[i])[-k:]) / k + for i in range(len(embeddings)) + ] + ) + + +def mean_rank_first_cross_category_neighbor( + embeddings: np.ndarray, metadata: List[Dict] +) -> float: + """Average rank when others are sorted by similarity: position of the closest other-category item. + + Rank 1 means the most similar product to you is already from another category. + Higher values mean other categories start further down the similarity-sorted list. + """ + n = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True) + sim = n @ n.T + np.fill_diagonal(sim, -np.inf) + cats = [m["category"] for m in metadata] + ranks: List[int] = [] + for i in range(len(embeddings)): + order = np.argsort(-sim[i]) + for rank_pos, j in enumerate(order, start=1): + if cats[j] != cats[i]: + ranks.append(rank_pos) + break + return float(np.mean(ranks)) if ranks else float("nan") + + +def mean_intra_inter_sim_gap(embeddings: np.ndarray, metadata: List[Dict]) -> float: + """Average gap: similarity to same-category products minus similarity to other categories.""" + n = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True) + sim = n @ n.T + np.fill_diagonal(sim, np.nan) + cats = [m["category"] for m in metadata] + gaps: List[float] = [] + for i in range(len(embeddings)): + intra = [ + sim[i, j] + for j in range(len(embeddings)) + if j != i and cats[j] == cats[i] + ] + inter = [sim[i, j] for j in range(len(embeddings)) if cats[j] != cats[i]] + if not intra or not inter: + continue + gaps.append(float(np.nanmean(intra) - np.nanmean(inter))) + return float(np.mean(gaps)) if gaps else float("nan") + + +def compute_category_precision_at_5(embeddings, metadata): + return compute_category_precision_at_k(embeddings, metadata, k=5) + + +def print_embedding_quality_report( + base_embs: np.ndarray, + finetuned_embs: np.ndarray, + metadata: List[Dict], +) -> None: + """Print a side-by-side base vs fine-tuned comparison on two robust metrics. + + We deliberately skip the easier "top-k neighbors share the category" score — + on this demo the product names already leak the category, so the base model + looks near-perfect there. The two checks below measure geometry in a way + that still shifts after contrastive training. + """ + rank_base = mean_rank_first_cross_category_neighbor(base_embs, metadata) + rank_ft = mean_rank_first_cross_category_neighbor(finetuned_embs, metadata) + gap_base = mean_intra_inter_sim_gap(base_embs, metadata) + gap_ft = mean_intra_inter_sim_gap(finetuned_embs, metadata) + + print("1) Sort other products by similarity (best match at the top).") + print(" How far down is the closest one that is not in your category?") + print(" (Higher = wrong-category items sit lower in that sorted list.)") + print(f" Base model: {rank_base:.1f}") + print(f" Fine-tuned: {rank_ft:.1f}") + print(f" Change: {rank_ft - rank_base:+.1f}") + print() + print("2) How much closer same-category items are than other categories") + print(" (larger = clearer gap between own category and the rest)") + print(f" Base model: {gap_base:.4f}") + print(f" Fine-tuned: {gap_ft:.4f}") + print(f" Change: {gap_ft - gap_base:+.4f}") diff --git a/tests/ecommerce_end_to_end/tests.sh b/tests/ecommerce_end_to_end/tests.sh new file mode 100644 index 000000000..b1d607d2a --- /dev/null +++ b/tests/ecommerce_end_to_end/tests.sh @@ -0,0 +1,4 @@ +#!/usr/bin/env bash +set -euo pipefail +pip install papermill +papermill notebook.ipynb output.ipynb -k python3 --log-output