feat: Add streamlit UI and categorical variable example

iusztinpaul · iusztinpaul · commit 6eff75cbe4ec · 2024-12-12T19:08:32.000+02:00
diff --git a/1_eda.ipynb b/1_eda.ipynb
@@ -22,7 +22,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -250,6 +250,34 @@
     "len(processed_df)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let's take a quick look at our categories:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "categories = set()\n",
+    "for category_list in processed_df[\"category\"]:\n",
+    "    categories.update(category_list)\n",
+    "categories"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "len(categories)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -280,7 +308,7 @@
    "outputs": [],
    "source": [
     "print(\"\\nNaN value counts:\")\n",
-    "print(processed_df[[\"price\", \"review_count\", \"review_rating\"]].isna().sum())"
+    "print(processed_df[[\"category\", \"price\", \"review_count\", \"review_rating\"]].isna().sum())"
    ]
   }
  ],
diff --git a/2_tabular_semantic_search_superlinked.ipynb b/2_tabular_semantic_search_superlinked.ipynb
diff --git a/3_tabular_semantic_search_text_to_sql.ipynb b/3_tabular_semantic_search_text_to_sql.ipynb
@@ -9,7 +9,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "\u001b[32m2024-12-11 18:15:51.463\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36msuperlinked_app.config\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m9\u001b[0m - \u001b[1mLoading '.env' file from: /Users/pauliusztin/Documents/01_projects/hands-on-retrieval/.env\u001b[0m\n"
+      "\u001b[32m2024-12-12 18:33:52.629\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36msuperlinked_app.config\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m9\u001b[0m - \u001b[1mLoading '.env' file from: /Users/pauliusztin/Documents/01_projects/hands-on-retrieval/.env\u001b[0m\n"
      ]
     }
    ],
@@ -258,7 +258,7 @@
     }
    ],
    "source": [
-    "df = pd.read_json(settings.PROCESSED_DATASET_PATH, lines=True)\n",
+    "df = pd.read_json(settings.PROCESSED_DATASET_PATH, lines=True).drop(columns=['category'])\n",
     "df.head()"
    ]
   },
@@ -378,7 +378,7 @@
     {
      "data": {
       "text/markdown": [
-       "**Node ID:** 92b1a07a-ac84-4da8-aab3-d562fe52edaf<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'All Aboard! New York: A City Primer', 'price': 9.99}<br>"
+       "**Node ID:** fe83bde8-54c6-46a7-bd31-7661bab4f8a8<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'All Aboard! New York: A City Primer', 'price': 9.99}<br>"
       ],
       "text/plain": [
        "<IPython.core.display.Markdown object>"
@@ -390,7 +390,7 @@
     {
      "data": {
       "text/markdown": [
-       "**Node ID:** 924225f0-acca-4e27-9edd-41534b4a54a2<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'Feminist Baby (Feminist Baby, 4)', 'price': 10.59}<br>"
+       "**Node ID:** 5ab49fa1-e254-4031-93bd-30dd0921c4ae<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'Feminist Baby (Feminist Baby, 4)', 'price': 10.59}<br>"
       ],
       "text/plain": [
        "<IPython.core.display.Markdown object>"
@@ -402,7 +402,7 @@
     {
      "data": {
       "text/markdown": [
-       "**Node ID:** 0c2201c5-8c0a-403a-875a-92903efc509f<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'The Mindful Dragon: A Dragon Book about Mindfulness. Teach Your Dragon To Be Mindful. A Cute Children Story to Teach Kids about Mindfulness, Focus and Peace. (My Dragon Books)', 'price': 11.69}<br>"
+       "**Node ID:** cc2142f6-58bf-44ee-9da5-e1df181c0d13<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'The Mindful Dragon: A Dragon Book about Mindfulness. Teach Your Dragon To Be Mindful. A Cute Children Story to Teach Kids about Mindfulness, Focus and Peace. (My Dragon Books)', 'price': 11.69}<br>"
       ],
       "text/plain": [
        "<IPython.core.display.Markdown object>"
@@ -428,7 +428,7 @@
     {
      "data": {
       "text/markdown": [
-       "**Node ID:** 21ef7ca7-72b0-407f-8f30-6b02b2b91cdc<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'Stables: Beautiful Paddocks, Horse Barns, and Tack Rooms', 'price': 53.1, 'review_rating': 4.7}<br>"
+       "**Node ID:** bbb01131-a9ca-4b8f-a06a-e5b713bbe0ff<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'Stables: Beautiful Paddocks, Horse Barns, and Tack Rooms', 'price': 53.1, 'review_rating': 4.7}<br>"
       ],
       "text/plain": [
        "<IPython.core.display.Markdown object>"
@@ -440,7 +440,7 @@
     {
      "data": {
       "text/markdown": [
-       "**Node ID:** e8f41cdf-af01-4134-9bf1-40a135be8ec4<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'The Mindful Dragon: A Dragon Book about Mindfulness. Teach Your Dragon To Be Mindful. A Cute Children Story to Teach Kids about Mindfulness, Focus and Peace. (My Dragon Books)', 'price': 11.69, 'review_rating': 4.7}<br>"
+       "**Node ID:** a0c8a2f9-595b-4fb4-b140-0a0bf68f436f<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'The Mindful Dragon: A Dragon Book about Mindfulness. Teach Your Dragon To Be Mindful. A Cute Children Story to Teach Kids about Mindfulness, Focus and Peace. (My Dragon Books)', 'price': 11.69, 'review_rating': 4.7}<br>"
       ],
       "text/plain": [
        "<IPython.core.display.Markdown object>"
@@ -452,7 +452,7 @@
     {
      "data": {
       "text/markdown": [
-       "**Node ID:** 95f86c04-2b9b-4145-b851-be4931b9b487<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'Build Your Running Body (A Total-Body Fitness Plan for All Distance Runners, from Milers to Ultramarathoners—Run Farther, Faster, and Injury-Free)', 'price': 13.49, 'review_rating': 4.7}<br>"
+       "**Node ID:** c1b7b9c7-1ceb-43ee-b7b9-b93990352772<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'Build Your Running Body (A Total-Body Fitness Plan for All Distance Runners, from Milers to Ultramarathoners—Run Farther, Faster, and Injury-Free)', 'price': 13.49, 'review_rating': 4.7}<br>"
       ],
       "text/plain": [
        "<IPython.core.display.Markdown object>"
@@ -478,7 +478,7 @@
     {
      "data": {
       "text/markdown": [
-       "**Node ID:** 15713c72-c930-4eb0-90b3-2ad617e3da6b<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'Stables: Beautiful Paddocks, Horse Barns, and Tack Rooms', 'review_count': 100, 'price': 53.1}<br>"
+       "**Node ID:** 14d164c1-9926-486e-8616-4e945a345174<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'Stables: Beautiful Paddocks, Horse Barns, and Tack Rooms', 'review_count': 100, 'price': 53.1}<br>"
       ],
       "text/plain": [
        "<IPython.core.display.Markdown object>"
@@ -490,7 +490,7 @@
     {
      "data": {
       "text/markdown": [
-       "**Node ID:** 285e8590-5b71-4bd9-b81a-ae6125a9bcf4<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'The Mindful Dragon: A Dragon Book about Mindfulness. Teach Your Dragon To Be Mindful. A Cute Children Story to Teach Kids about Mindfulness, Focus and Peace. (My Dragon Books)', 'review_count': 623, 'price': 11.69}<br>"
+       "**Node ID:** b9a67ec0-7101-446d-b36c-bb8435442123<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'The Mindful Dragon: A Dragon Book about Mindfulness. Teach Your Dragon To Be Mindful. A Cute Children Story to Teach Kids about Mindfulness, Focus and Peace. (My Dragon Books)', 'review_count': 623, 'price': 11.69}<br>"
       ],
       "text/plain": [
        "<IPython.core.display.Markdown object>"
@@ -502,7 +502,7 @@
     {
      "data": {
       "text/markdown": [
-       "**Node ID:** 211e96fb-f435-49d3-a952-41209d479280<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'Build Your Running Body (A Total-Body Fitness Plan for All Distance Runners, from Milers to Ultramarathoners—Run Farther, Faster, and Injury-Free)', 'review_count': 573, 'price': 13.49}<br>"
+       "**Node ID:** 094e44be-bf37-4d56-9d9d-3045ffbc6004<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'Build Your Running Body (A Total-Body Fitness Plan for All Distance Runners, from Milers to Ultramarathoners—Run Farther, Faster, and Injury-Free)', 'review_count': 573, 'price': 13.49}<br>"
       ],
       "text/plain": [
        "<IPython.core.display.Markdown object>"
@@ -514,7 +514,7 @@
     {
      "data": {
       "text/markdown": [
-       "**Node ID:** 799f5a70-47c3-40af-b96a-be80fc1ccfa1<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'All Aboard! New York: A City Primer', 'review_count': 74, 'price': 9.99}<br>"
+       "**Node ID:** 2f8c362a-f7ea-4643-8986-107d0ba7ea5b<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'All Aboard! New York: A City Primer', 'review_count': 74, 'price': 9.99}<br>"
       ],
       "text/plain": [
        "<IPython.core.display.Markdown object>"
@@ -526,7 +526,7 @@
     {
      "data": {
       "text/markdown": [
-       "**Node ID:** 38440887-0a63-4915-abe3-8fad6fd0ad10<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'The Rise: Creativity, the Gift of Failure, and the Search for Mastery', 'review_count': 204, 'price': 13.79}<br>"
+       "**Node ID:** a3c0a9e4-c3b5-482e-853f-d66ac38f7787<br>**Similarity:** None<br>**Text:** <br>**Metadata:** {'title': 'The Rise: Creativity, the Gift of Failure, and the Search for Mastery', 'review_count': 204, 'price': 13.79}<br>"
       ],
       "text/plain": [
        "<IPython.core.display.Markdown object>"
diff --git a/INSTALL_AND_USAGE.md b/INSTALL_AND_USAGE.md
@@ -90,8 +90,9 @@ data/
 
 1. Start it up:
    ```bash
-   make start-server
+   make start-superlinked-server
    ```
+   FastAPI endpoints docs available at `http://localhost:8080/docs`
 
 2. Load your data:
    ```bash
@@ -104,4 +105,10 @@ data/
    make post-semantic-query   # Natural language search
    ```
 
+4. Start the Streamlit UI:
+   ```bash
+   make start-ui
+   ```
+   Accessible at `http://localhost:8501/`
+
 > 🔔 Make sure the server is running (step 1) before executing the data loading or search commands.
diff --git a/Makefile b/Makefile
@@ -31,4 +31,7 @@ post-semantic-query:
 	'http://localhost:8080/api/v1/search/semantic_query' \
 	-H 'accept: application/json' \
 	-H 'Content-Type: application/json' \
-	-d '{"natural_query": "books with a price lower than 100", "limit": 3}' | jq '.'
+	-d '{"natural_query": "books with a price lower than 100", "limit": 3}' | jq '.'
+
+start-ui:
+	uv run streamlit run tools/streamlit_app.py
diff --git a/pyproject.toml b/pyproject.toml
@@ -18,6 +18,7 @@ dependencies = [
     "zstandard>=0.23.0",
     "sqlalchemy>=2.0.36",
     "matplotlib>=3.9.3",
+    "streamlit>=1.41.0",
 ]
 
 [dependency-groups]
diff --git a/superlinked_app/constants.py b/superlinked_app/constants.py
@@ -0,0 +1,71 @@
+TYPES = ["product", "book"]
+
+CATEGORIES = [
+    "Accessories",
+    "Appliances",
+    "Arts & Photography",
+    "Arts, Crafts & Sewing",
+    "Automotive",
+    "Baby Care",
+    "Baby Products",
+    "Bath",
+    "Beauty & Personal Care",
+    "Bedding",
+    "Beverages",
+    "Biographies & Memoirs",
+    "Books",
+    "CDs & Vinyl",
+    "Camera & Photo",
+    "Cell Phones & Accessories",
+    "Children's Books",
+    "Christian Books & Bibles",
+    "Classical",
+    "Clothing, Shoes & Jewelry",
+    "Computers & Accessories",
+    "Costumes & Accessories",
+    "Dogs",
+    "Electrical",
+    "Electronics",
+    "Event & Party Supplies",
+    "Exercise & Fitness",
+    "Exterior Accessories",
+    "GPS, Finders & Accessories",
+    "Grocery & Gourmet Food",
+    "Hair Care",
+    "Health & Household",
+    "Home & Kitchen",
+    "Hunting & Fishing",
+    "Industrial & Scientific",
+    "Industrial Electrical",
+    "Kitchen & Dining",
+    "Lighting Assemblies & Accessories",
+    "Lights & Lighting Accessories",
+    "Luggage & Travel Gear",
+    "Makeup",
+    "Medical Supplies & Equipment",
+    "Men",
+    "Movies & TV",
+    "Musical Instruments",
+    "Office & School Supplies",
+    "Office Products",
+    "Patio Furniture & Accessories",
+    "Patio, Lawn & Garden",
+    "Pet Supplies",
+    "Pop",
+    "Portable Audio & Video",
+    "Power & Hand Tools",
+    "Raw Materials",
+    "Replacement Parts",
+    "Self-Help",
+    "Sports & Outdoor Play",
+    "Sports & Outdoors",
+    "Stuffed Animals & Plush Toys",
+    "Tires & Wheels",
+    "Tools & Home Improvement",
+    "Toy Figures & Playsets",
+    "Toys & Games",
+    "Vehicles",
+    "Video Games",
+    "Wall Art",
+    "Women",
+]
diff --git a/superlinked_app/data_processing.py b/superlinked_app/data_processing.py
@@ -1,8 +1,30 @@
+import random
 from typing import Optional
 
 import pandas as pd
 
 
+def parse_category(category: str) -> list[str]:
+    """Parse the category string and return the first category.
+
+    Args:
+        category: String containing the category list (e.g., "['Books', 'Fiction', 'Literature']")
+
+    Returns:
+        String containing the first category, or None if parsing fails
+    """
+    if isinstance(category, list) and pd.isna(category).any():
+        return []
+    elif not isinstance(category, list) and pd.isna(category):
+        return []
+
+    try:
+        keep_num_categories = 1 if random.random() < 0.9 else 2
+        return [c.strip() for c in category][:keep_num_categories]
+    except (ValueError, IndexError):
+        return []
+
+
 def parse_stars(stars: str) -> Optional[float]:
     """Parse the stars rating from a string to a float value between 0 and 5.
 
@@ -89,13 +111,16 @@ def process_amazon_dataset(df: pd.DataFrame) -> pd.DataFrame:
             - price (float): Price value
     """
 
+    random.seed(6)
+
     # Create a copy to avoid modifying the original DataFrame
     df_processed = df.copy()
 
     # Keep only required columns
     columns_to_keep = [
         "asin",
         "type",
+        "category",
         "title",
         "description",
         "stars",
@@ -105,6 +130,7 @@ def process_amazon_dataset(df: pd.DataFrame) -> pd.DataFrame:
     df_processed = df_processed[columns_to_keep]
 
     # Apply transformations
+    df_processed["category"] = df_processed["category"].apply(parse_category)
     df_processed["review_rating"] = df_processed["stars"].apply(parse_stars)
     df_processed["review_count"] = df_processed["ratings"].apply(parse_ratings)
     df_processed["price"] = df_processed["price"].apply(parse_price)
diff --git a/superlinked_app/index.py b/superlinked_app/index.py
@@ -1,9 +1,12 @@
 from superlinked import framework as sl
 
+from superlinked_app import constants
+
 
 class ProductSchema(sl.Schema):
     id: sl.IdField
     type: sl.String
+    category: sl.StringList
     title: sl.String
     description: sl.String
     review_rating: sl.Float
@@ -13,6 +16,12 @@ class ProductSchema(sl.Schema):
 
 product = ProductSchema()
 
+category_space = sl.CategoricalSimilaritySpace(
+    category_input=product.category,
+    categories=constants.CATEGORIES,
+    uncategorized_as_category=True,
+    negative_filter=-1,
+)
 description_space = sl.TextSimilaritySpace(
     text=product.description, model="Alibaba-NLP/gte-large-en-v1.5"
 )
@@ -24,6 +33,11 @@ class ProductSchema(sl.Schema):
 )
 
 product_index = sl.Index(
-    spaces=[description_space, review_rating_maximizer_space, price_minimizer_space],
-    fields=[product.type, product.review_rating, product.price],
+    spaces=[
+        category_space,
+        description_space,
+        review_rating_maximizer_space,
+        price_minimizer_space,
+    ],
+    fields=[product.type, product.category, product.review_rating, product.price],
 )
diff --git a/superlinked_app/query.py b/superlinked_app/query.py
diff --git a/tools/streamlit_app.py b/tools/streamlit_app.py
diff --git a/uv.lock b/uv.lock

Original file line number	Diff line number	Diff line change
`@@ -18,6 +18,7 @@ dependencies = [`
`18`	`18`	`"zstandard>=0.23.0",`
`19`	`19`	`"sqlalchemy>=2.0.36",`
`20`	`20`	`"matplotlib>=3.9.3",`
	`21`	`+ "streamlit>=1.41.0",`
`21`	`22`	`]`
`22`	`23`
`23`	`24`	`[dependency-groups]`