Spaces:
Runtime error
Runtime error
| 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()) | |