import os import re import numpy as np import pandas as pd import gradio as gr import torch from sentence_transformers import SentenceTransformer import faiss from transformers import AutoTokenizer, AutoModelForSeq2SeqLM # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- DATA_PATH = os.path.join(os.path.dirname(__file__), "products.csv") EMBED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2" GEN_MODEL_NAME = "google/flan-t5-base" # swap to "google/flan-t5-small" for faster/lower-memory CPU TOP_K = 5 # --------------------------------------------------------------------------- # Load data # --------------------------------------------------------------------------- print("Loading product catalog...") df = pd.read_csv(DATA_PATH) df = df.fillna("") def row_to_text(row): return ( f"{row['name']}. Category: {row['category']} > {row['subcategory']}. " f"Brand: {row['brand']}. Price: ${row['price_usd']}. " f"Rating: {row['rating']} ({row['num_reviews']} reviews). " f"Color: {row['color']}. Material: {row['material']}. " f"Tags: {row['tags']}. Description: {row['description']}" ) df["embedding_text"] = df.apply(row_to_text, axis=1) # --------------------------------------------------------------------------- # Build (or load cached) embeddings + FAISS index # --------------------------------------------------------------------------- print("Loading embedding model...") embedder = SentenceTransformer(EMBED_MODEL_NAME) EMB_CACHE = os.path.join(os.path.dirname(__file__), "product_embeddings.npy") if os.path.exists(EMB_CACHE): print("Loading cached embeddings...") embeddings = np.load(EMB_CACHE) if embeddings.shape[0] != len(df): embeddings = None else: embeddings = None if embeddings is None: print("Computing embeddings for the catalog (this happens once)...") embeddings = embedder.encode( df["embedding_text"].tolist(), batch_size=64, show_progress_bar=True, convert_to_numpy=True, normalize_embeddings=True, ).astype("float32") np.save(EMB_CACHE, embeddings) dim = embeddings.shape[1] index = faiss.IndexFlatIP(dim) # inner product on normalized vectors = cosine similarity index.add(embeddings) print(f"Index ready with {index.ntotal} products.") # --------------------------------------------------------------------------- # Load small generator model # --------------------------------------------------------------------------- print("Loading generation model (flan-t5-base)... this may take a minute on first run.") gen_tokenizer = AutoTokenizer.from_pretrained(GEN_MODEL_NAME) gen_model = AutoModelForSeq2SeqLM.from_pretrained(GEN_MODEL_NAME) def generate_text(prompt, max_new_tokens=200): """Direct generate() call instead of pipeline() — avoids relying on the transformers pipeline task registry, whose task names have changed across major transformers versions (e.g. 'text2text-generation' was removed as a pipeline alias in transformers 5.x). Runs on CPU.""" inputs = gen_tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024) output_ids = gen_model.generate(**inputs, max_new_tokens=max_new_tokens) return gen_tokenizer.decode(output_ids[0], skip_special_tokens=True).strip() # --------------------------------------------------------------------------- # Helpers: lightweight query understanding # --------------------------------------------------------------------------- PRICE_UNDER_RE = re.compile(r"(?:under|below|less than|cheaper than|<)\s*\$?\s*(\d+(?:\.\d+)?)", re.I) PRICE_OVER_RE = re.compile(r"(?:over|above|more than|>)\s*\$?\s*(\d+(?:\.\d+)?)", re.I) CATEGORIES = df["category"].unique().tolist() def extract_constraints(query): max_price = None min_price = None m = PRICE_UNDER_RE.search(query) if m: max_price = float(m.group(1)) m = PRICE_OVER_RE.search(query) if m: min_price = float(m.group(1)) mentioned_category = None q_lower = query.lower() for cat in CATEGORIES: if cat.lower() in q_lower or any(word.lower() in q_lower for word in cat.split(" & ")): mentioned_category = cat break return {"max_price": max_price, "min_price": min_price, "category": mentioned_category} def retrieve(query, top_k=TOP_K): constraints = extract_constraints(query) q_emb = embedder.encode([query], convert_to_numpy=True, normalize_embeddings=True).astype("float32") # Over-fetch, then filter/re-rank so price/category constraints don't starve results scores, idxs = index.search(q_emb, min(60, len(df))) idxs = idxs[0] scores = scores[0] candidates = df.iloc[idxs].copy() candidates["score"] = scores if constraints["max_price"] is not None: candidates = candidates[candidates["price_usd"] <= constraints["max_price"]] if constraints["min_price"] is not None: candidates = candidates[candidates["price_usd"] >= constraints["min_price"]] if constraints["category"] is not None: filtered = candidates[candidates["category"] == constraints["category"]] if len(filtered) > 0: candidates = filtered if len(candidates) == 0: # fall back to unfiltered top results if constraints were too strict candidates = df.iloc[idxs].copy() candidates["score"] = scores return candidates.head(top_k) def format_context(products): lines = [] for _, p in products.iterrows(): lines.append( f"- {p['name']} | {p['category']} / {p['subcategory']} | ${p['price_usd']:.2f} | " f"{p['rating']}★ ({p['num_reviews']} reviews) | {p['stock_status']} | {p['description']}" ) return "\n".join(lines) def format_sources_markdown(products): if len(products) == 0: return "_No matching products found in the catalog._" lines = ["**Retrieved from catalog:**\n"] for _, p in products.iterrows(): lines.append( f"- **{p['name']}** — ${p['price_usd']:.2f} · {p['rating']}★ " f"({int(p['num_reviews'])} reviews) · {p['stock_status']} · `{p['product_id']}`" ) return "\n".join(lines) def build_prompt(query, context): return ( "You are a helpful e-commerce shopping assistant. Using ONLY the product " "information below, answer the customer's question. Recommend specific " "products by name, mention price, and be concise and friendly. If nothing " "in the catalog fits, say so honestly.\n\n" f"Product catalog matches:\n{context}\n\n" f"Customer question: {query}\n\n" "Answer:" ) def rag_answer(message, history): products = retrieve(message, top_k=TOP_K) context = format_context(products) prompt = build_prompt(message, context) generated = generate_text(prompt) if not generated: generated = "I couldn't generate a response — please try rephrasing your question." sources_md = format_sources_markdown(products) full_reply = f"{generated}\n\n---\n{sources_md}" return full_reply # --------------------------------------------------------------------------- # Gradio UI # --------------------------------------------------------------------------- EXAMPLES = [ "I need wireless headphones under $100", "Recommend a good gift for someone who loves cooking", "What running shoes do you have over $80?", "Show me eco-friendly toys for kids", "I want a durable yoga mat, budget friendly", ] with gr.Blocks(title="RAG E-commerce Assistant") as demo: gr.Markdown( """ # 🛍️ RAG E-commerce Shopping Assistant Ask about products in natural language — this assistant retrieves real matches from a **2,100+ item product catalog** and uses them to ground its answer (Retrieval-Augmented Generation). """ ) chatbot = gr.ChatInterface( fn=rag_answer, examples=EXAMPLES, chatbot=gr.Chatbot(height=480), textbox=gr.Textbox(placeholder="e.g. I need a laptop bag under $50", scale=7), ) gr.Markdown( "Powered by `sentence-transformers/all-MiniLM-L6-v2` (retrieval) + " "`google/flan-t5-base` (generation) + FAISS. Catalog: `products.csv`." ) if __name__ == "__main__": demo.launch(theme=gr.themes.Soft())