Update smart_suggestion/flan_suggestor.py

smart_suggestion/flan_suggestor.py

@@ -1,71 +1,74 @@
-import os
-import torch
-import pandas as pd
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-
-# Load model
-model_name = "google/flan-t5-small"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-
-# Load all product CSVs from nested directories
-def load_all_product_data(base_path="data"):
-    all_data = []
-    for root, dirs, files in os.walk(base_path):
-        for file in files:
-            if file.endswith(".csv"):
-                full_path = os.path.join(root, file)
-                df = pd.read_csv(full_path)
-                df["Brand"] = os.path.splitext(file)[0]
-                df["Store"] = root.split(os.sep)[-4]  # e.g., "Store C"
-                df["Category"] = root.split(os.sep)[-2]
-                all_data.append(df)
-    return pd.concat(all_data, ignore_index=True)
-
-df = load_all_product_data("data")
-
-# Generate smart suggestion
-def generate_product_description(prompt):
-    prompt = prompt.lower()
-
-    # Basic price filter
-    price_limit = 99999
-    if "under" in prompt:
-        try:
-            price_limit = int(prompt.split("under")[-1].split()[0])
-        except:
-            pass
-
-    filtered_df = df[df["Price"] <= price_limit]
-    filtered_df = filtered_df[df["In Stock"].str.lower() == "yes"]
-
-    if "dry hair" in prompt:
-        filtered_df = filtered_df[filtered_df["Hair Type"].str.lower().str.contains("dry", na=False)]
-    elif "oily hair" in prompt:
-        filtered_df = filtered_df[filtered_df["Hair Type"].str.lower().str.contains("oily", na=False)]
-    elif "normal hair" in prompt:
-        filtered_df = filtered_df[filtered_df["Hair Type"].str.lower().str.contains("normal", na=False)]
-
-    if "gift" in prompt:
-        filtered_df = filtered_df[filtered_df["Tags"].str.contains("gift", case=False, na=False)]
-    if "budget" in prompt:
-        filtered_df = filtered_df[filtered_df["Tags"].str.contains("budget", case=False, na=False)]
-
-    if filtered_df.empty:
-        return "🤷 Sorry, no matching suggestions found."
-
-    rows = []
-    for _, row in filtered_df.iterrows():
-        text = f"{row['Brand']} {row['Quantity']} – ₹{row['Price']} (Floor {row['Floor']}, Aisle {row['Aisle']})"
-        if pd.notna(row.get("Offer")) and str(row["Offer"]).strip():
-            text += f" | 🎉 {row['Offer']}"
-        rows.append(text)
-
-    product_text = "\n".join(rows)
-    model_prompt = f"Suggest top products:\n{product_text}"
-
-    input_ids = tokenizer(model_prompt, return_tensors="pt").input_ids
-    with torch.no_grad():
-        output_ids = model.generate(input_ids, max_new_tokens=100)
-
-    return tokenizer.decode(output_ids[0], skip_special_tokens=True)
+import os
+import torch
+import pandas as pd
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+# Load model
+model_name = "google/flan-t5-small"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+# Load all product CSVs from nested directories
+def load_all_product_data(base_path="data"):
+    all_data = []
+    for root, dirs, files in os.walk(base_path):
+        for file in files:
+            if file.endswith(".csv"):
+                full_path = os.path.join(root, file)
+                df = pd.read_csv(full_path)
+                df["Brand"] = os.path.splitext(file)[0]
+                df["Store"] = root.split(os.sep)[-4]  # e.g., "Store C"
+                df["Category"] = root.split(os.sep)[-2]
+                all_data.append(df)
+    return pd.concat(all_data, ignore_index=True)
+
+df = load_all_product_data("data")
+
+# Generate smart suggestion
+def generate_product_description(prompt):
+    prompt = prompt.lower()
+
+    # Basic price filter
+    price_limit = 99999
+    if "under" in prompt:
+        try:
+            price_limit = int(prompt.split("under")[-1].split()[0])
+        except:
+            pass
+
+    filtered_df = df[df["Price"] <= price_limit]
+    filtered_df = filtered_df[df["In Stock"].str.lower() == "yes"]
+
+    if "dry hair" in prompt:
+        filtered_df = filtered_df[filtered_df["Hair Type"].str.lower().str.contains("dry", na=False)]
+    elif "oily hair" in prompt:
+        filtered_df = filtered_df[filtered_df["Hair Type"].str.lower().str.contains("oily", na=False)]
+    elif "normal hair" in prompt:
+        filtered_df = filtered_df[filtered_df["Hair Type"].str.lower().str.contains("normal", na=False)]
+
+    if "gift" in prompt:
+        filtered_df = filtered_df[filtered_df["Tags"].str.contains("gift", case=False, na=False)]
+    if "budget" in prompt:
+        filtered_df = filtered_df[filtered_df["Tags"].str.contains("budget", case=False, na=False)]
+
+    if filtered_df.empty:
+        return "🤷 Sorry, no matching suggestions found."
+
+    rows = []
+    for _, row in filtered_df.iterrows():
+        line = f"{row['Brand']} {row['Quantity']} – ₹{row['Price']} (Floor {row['Floor']}, Aisle {row['Aisle']})"
+        if pd.notna(row.get("Offer")) and str(row["Offer"]).strip():
+            line += f"\n🎉 {row['Offer']}"
+        rows.append(line)
+
+    product_text = "\n".join(rows)
+    model_prompt = f"Suggest top products:\n{product_text}"
+
+    input_ids = tokenizer(model_prompt, return_tensors="pt").input_ids
+    with torch.no_grad():
+        output_ids = model.generate(input_ids, max_new_tokens=100)
+
+    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+    # Combine product suggestions and model response line by line
+    return "\n\n".join(rows[:5]) + "\n\n🧠 AI Suggestion: " + response