Add data prep and vector DB

requirements.txt

@@ -1,6 +1,7 @@
 gradio==5.42.0
 openai
 python-dotenv==1.0.0
+datasets
 numpy
 pandas
 scikit-learn

src/modules/data_prep.py

@@ -0,0 +1,147 @@
+import pandas as pd
+from datasets import load_dataset
+from pathlib import Path
+import numpy as np
+import faiss
+from openai import OpenAI
+from dotenv import load_dotenv
+import os
+from src.config import EMBEDDING_MODEL
+
+load_dotenv()
+
+_FILE_PATH = Path(__file__).parents[2]
+
+
+def load_amazon_raw_product_data() -> pd.DataFrame:
+    ds = load_dataset("ckandemir/amazon-products")
+    df = ds["train"].to_pandas()
+    return df
+
+
+def load_clean_amazon_product_data() -> pd.DataFrame:
+    return pd.read_parquet(_FILE_PATH / "data" / "amazon_products.parquet")
+
+
+def prepare_amazon_product_data(df: pd.DataFrame) -> pd.DataFrame:
+    """
+    Data preparation for Amazon products.
+
+    Args:
+        df: DataFrame with 'Product Name', 'Category', 'Description' columns
+
+    Returns:
+        DataFrame
+    """
+    # Full text is combination of Category + Description
+    df.loc[:, "FullText"] = (
+        df["Product Name"] + " | " + df["Category"] + " | " + df["Description"]
+    )
+    df.loc[:, "FullText"] = df.FullText.str.lower().str.strip().str.replace("\n", " ")
+
+    df[["MainCategory", "SecondaryCategory", "TertiaryCategory"]] = df[
+        "Category"
+    ].str.split(r" \| ", n=2, expand=True, regex=True)
+    df = df.dropna(subset=["MainCategory", "SecondaryCategory"])
+
+    # Drop dupes
+    df = df.drop_duplicates(subset=["FullText"])
+
+    return df.loc[
+        :,
+        [
+            "Product Name",
+            "Description",
+            "MainCategory",
+            "SecondaryCategory",
+            "TertiaryCategory",
+            "FullText",
+        ],
+    ]
+
+
+def save_as_parquet(df: pd.DataFrame):
+    """
+    Save DataFrame to parquet file.
+    """
+    df.to_parquet(_FILE_PATH / "data" / "amazon_products.parquet", index=False)
+    print(f"Saved to {_FILE_PATH / 'data' / 'amazon_products.parquet'}")
+
+
+def create_faiss_index(df: pd.DataFrame, batch_size: int = 100):
+    """
+    Create FAISS index from product data using Fireworks AI embeddings.
+
+    Args:
+        df: DataFrame with 'FullText' column to embed
+        batch_size: Number of texts to embed in each API call
+
+    Returns:
+        Tuple of (faiss_index, embeddings_array)
+    """
+    assert (
+        os.getenv("FIREWORKS_API_KEY") is not None
+    ), "FIREWORKS_API_KEY not found in environment variables"
+    client = OpenAI(
+        api_key=os.getenv("FIREWORKS_API_KEY"),
+        base_url="https://api.fireworks.ai/inference/v1",
+    )
+
+    print(f"Generating embeddings for {len(df)} products...")
+
+    all_embeddings = []
+    texts = df["FullText"].tolist()
+
+    for i in range(0, len(texts), batch_size):
+        batch = texts[i : i + batch_size]
+        print(
+            f"Processing batch {i // batch_size + 1}/{(len(texts) + batch_size - 1) // batch_size}"
+        )
+
+        response = client.embeddings.create(model=EMBEDDING_MODEL, input=batch)
+
+        batch_embeddings = [item.embedding for item in response.data]
+        all_embeddings.extend(batch_embeddings)
+
+    embeddings_array = np.array(all_embeddings, dtype=np.float32)
+
+    dimension = embeddings_array.shape[1]
+    index = faiss.IndexFlatL2(
+        dimension
+    )  # L2 distance for cosine similarity after normalization
+
+    # Normalize embeddings for cosine similarity
+    faiss.normalize_L2(embeddings_array)
+
+    index.add(embeddings_array)
+
+    print(f"Created FAISS index with {index.ntotal} vectors of dimension {dimension}")
+
+    faiss.write_index(index, str(_FILE_PATH / "data" / "faiss_index.bin"))
+    np.save(_FILE_PATH / "data" / "embeddings.npy", embeddings_array)
+
+    print(f"Saved FAISS index to {_FILE_PATH / 'data' / 'faiss_index.bin'}")
+    print(f"Saved embeddings to {_FILE_PATH / 'data' / 'embeddings.npy'}")
+
+    return index, embeddings_array
+
+
+def load_faiss_index():
+    """
+    Load pre-computed FAISS index and embeddings from disk.
+
+    Returns:
+        Tuple of (faiss_index, embeddings_array)
+    """
+    index = faiss.read_index(str(_FILE_PATH / "data" / "faiss_index.bin"))
+    embeddings = np.load(_FILE_PATH / "data" / "embeddings.npy")
+    print(f"Loaded FAISS index with {index.ntotal} vectors")
+    return index, embeddings
+
+
+if __name__ == "__main__":
+    _df = load_amazon_raw_product_data()
+    _df = prepare_amazon_product_data(_df)
+    save_as_parquet(_df)
+
+    create_faiss_index(_df)