Filter to only top 5 categories

requirements.txt

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

src/data_prep/data_prep.py

@@ -44,6 +44,19 @@ def prepare_amazon_product_data(df: pd.DataFrame) -> pd.DataFrame:
     # Drop dupes
     df = df.drop_duplicates(subset=["FullText"])
 
+    # Downsample where MainCategory == Toys and Games to 650 since in raw data its over 70% of data
+    df_non_toys = df[df["MainCategory"] != "Toys & Games"]
+    df_toys = df[df["MainCategory"] == "Toys & Games"]
+    df_toys = df_toys.sample(n=650, random_state=42)
+    df = pd.concat([df_non_toys, df_toys])
+
+    # Filter to only top 5 MainCategories
+    df = df[df["MainCategory"].isin(df["MainCategory"].value_counts().index[:5])]
+
+    print(
+        f"Prepared dataset with {len(df)} products with \n Count of MainCategories: {df['MainCategory'].value_counts()}"
+    )
+
     return df.loc[
         :,
         [