import json import html import pandas as pd import gradio as gr import torch from transformers import AutoModelForCausalLM, AutoTokenizer MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct" ORDER_COLUMNS = [ "customer", "item", "quantity", "flavor", "pickup_time", "delivery_address", "payment_status", "notes", "missing_fields", ] tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32) model.eval() SYSTEM_PROMPT = """ You are a careful order extraction engine for tiny sellers. Extract customer orders from messy DMs. Return only valid JSON with this exact shape: { "orders": [ { "customer": "", "item": "", "quantity": "", "flavor": "", "pickup_time": "", "delivery_address": "", "payment_status": "", "notes": "", "missing_fields": [] } ], "prep_list": [], "reply_drafts": [] } Critical rules: - Treat each line as one separate customer message. - The text before the first ":" is the customer name. - Copy customer names exactly as written. Do not uppercase or lowercase them. - Never copy details from one customer's message into another customer's order. - Include every customer message that looks like an order or possible order. - Use only facts explicitly present in that customer's own message. - If a value is unknown, use an empty string. - Do not add order_id or total_cost. - For pickup orders, put pickup time in pickup_time. Put a pickup place or delivery address in delivery_address. - If the customer is unsure, still include the order and describe the uncertainty in notes. - missing_fields should only include fields the seller needs to ask for: quantity, flavor, pickup_time, delivery_address, payment_status. - Always set prep_list to []. - Always set reply_drafts to []. """ SINGLE_ORDER_PROMPT = """ You extract one order from one customer's DM. Return only valid JSON with this exact shape: { "item": "", "quantity": "", "flavor": "", "pickup_time": "", "delivery_address": "", "payment_status": "", "notes": "", "missing_fields": [] } Rules: - Use only facts from this one message. - Do not invent details. - Put dates and times in pickup_time, such as "tomorrow", "Saturday morning", or "Friday 5pm". - Put pickup places or delivery addresses in delivery_address, such as "farmers market". - "pickup at the farmers market" means delivery_address is "farmers market", not pickup_time. - "paid already" means payment_status is "paid". - "I can pay Venmo" means payment_status is "can pay Venmo". - If unknown, use an empty string. - Do not ask for flavor unless the product clearly needs a flavor choice. - missing_fields can only contain: quantity, flavor, pickup_time, delivery_address, payment_status. """ EXAMPLE_INPUT = """Maya: Hi! Can I get 2 dozen cupcakes for Saturday morning? Half vanilla, half chocolate. Sam: Need 1 birthday cake, chocolate, for pickup Friday 5pm. I can pay Venmo. Lena: Do you still have lemon bars? I need some for tomorrow but not sure how many yet. Chris: 12 cookies please, pickup at the farmers market. Paid already. """ FOOD_TRUCK_EXAMPLE = """Alex: Can I get 3 chicken tacos for pickup at 6:30 tonight? Paid on Cash App. Jamie: Do you still have vegan bowls? Need 2 tomorrow for office lunch. Priya: One brisket sandwich, no onions. I'll pick up at the truck on Main Street. Nate: 4 lemonades for the soccer team, pickup after practice. """ CRAFT_SELLER_EXAMPLE = """Olivia: I want 2 custom mugs with blue initials. Can you ship to 18 Pine Road? Ben: Need one candle gift box for Saturday. Lavender if you have it. Rosa: Can I order 3 tote bags? I can pick up at the market. Eli: Do you still make birthday stickers? Need some next week but not sure how many. """ FARMERS_MARKET_EXAMPLE = """Grace: Can you hold 2 sourdough loaves for Sunday pickup? Leo: I need 1 jar of strawberry jam and 2 honey bottles. Paid already. Mina: Do you have eggs this weekend? Maybe 2 dozen if available. Noah: Please save me 3 bags of granola, pickup at the farmers market. """ EXAMPLES = [ [EXAMPLE_INPUT], [FOOD_TRUCK_EXAMPLE], [CRAFT_SELLER_EXAMPLE], [FARMERS_MARKET_EXAMPLE], ] def extract_json(text): start = text.find("{") end = text.rfind("}") if start == -1 or end == -1: raise ValueError("No JSON object found") return json.loads(text[start:end + 1]) def normalize_orders(data): rows = [] for order in data.get("orders", []): row = {} for col in ORDER_COLUMNS: value = order.get(col, "") if isinstance(value, list): value = ", ".join(str(v) for v in value) row[col] = value rows.append(row) return pd.DataFrame(rows, columns=ORDER_COLUMNS) def format_list(title, items): if not items: return f"### {title}\nNothing found." lines = [] for item in items: if isinstance(item, dict): lines.append("- " + json.dumps(item, ensure_ascii=False)) else: lines.append(f"- {item}") return f"### {title}\n" + "\n".join(lines) def format_replies(replies): if not replies: return "### Reply drafts\nNothing found." lines = [] for reply in replies: customer = reply.get("customer", "Customer") text = reply.get("reply", "") lines.append(f"**{customer}**\n\n{text}") return "### Reply drafts\n\n" + "\n\n---\n\n".join(lines) def text_value(value): if isinstance(value, list): return ", ".join(str(v) for v in value if str(v).strip()) if value is None: return "" return str(value).strip() def missing_list(order): raw = order.get("missing_fields", []) if isinstance(raw, str): fields = [part.strip() for part in raw.split(",") if part.strip()] else: fields = [str(part).strip() for part in raw if str(part).strip()] allowed = {"quantity", "flavor", "pickup_time", "delivery_address", "payment_status"} fields = [field for field in fields if field in allowed] item = text_value(order.get("item")) quantity = text_value(order.get("quantity")) flavor = text_value(order.get("flavor")) pickup_time = text_value(order.get("pickup_time")) delivery_address = text_value(order.get("delivery_address")) payment_status = text_value(order.get("payment_status")) if item and not quantity: fields.append("quantity") if pickup_time: fields = [field for field in fields if field != "pickup_time"] if delivery_address: fields = [field for field in fields if field != "delivery_address"] if payment_status: fields = [field for field in fields if field != "payment_status"] if flavor: fields = [field for field in fields if field != "flavor"] if "flavor" in fields and item.lower() not in ["cake", "birthday cake", "cupcakes"]: fields = [field for field in fields if field != "flavor"] return sorted(set(fields)) def post_process_order(order, message): msg = message.lower() if "paid already" in msg or "already paid" in msg: order["payment_status"] = "paid" elif "venmo" in msg: order["payment_status"] = "can pay Venmo" elif "paid" not in msg and "venmo" not in msg: order["payment_status"] = "" pickup_time = text_value(order.get("pickup_time")) if "paid" in pickup_time.lower() or "venmo" in pickup_time.lower(): order["pickup_time"] = "" if "farmers market" in msg: order["delivery_address"] = "farmers market" if "farmers market" in text_value(order.get("pickup_time")).lower(): order["pickup_time"] = "" order["missing_fields"] = missing_list(order) return order def build_prep_list(data): items = [] for order in data.get("orders", []): item = text_value(order.get("item")) if not item: continue customer = text_value(order.get("customer")) or "customer" quantity = text_value(order.get("quantity")) or "quantity to confirm" flavor = text_value(order.get("flavor")) line = f"{quantity} {item}" if flavor: line += f" ({flavor})" line += f" - {customer}" items.append(line) return items def build_reply_drafts(data): replies = [] labels = { "quantity": "quantity", "flavor": "flavor", "pickup_time": "pickup or delivery time", "delivery_address": "pickup place", "payment_status": "payment status", } for order in data.get("orders", []): customer = text_value(order.get("customer")) or "there" item = text_value(order.get("item")) or "order" quantity = text_value(order.get("quantity")) flavor = text_value(order.get("flavor")) missing = [labels.get(field, field) for field in missing_list(order)] if missing: needed = ", ".join(missing) reply = f"Thanks, {customer}! I have your {item} order. Could you confirm the {needed}?" else: summary = f"{quantity} {item}".strip() if flavor: summary += f" ({flavor})" reply = f"Thanks, {customer}! Confirming your order: {summary}." replies.append({"customer": customer, "reply": reply}) return replies def build_summary(data): orders = data.get("orders", []) total = len(orders) followups = sum(1 for order in orders if missing_list(order)) ready = total - followups return f"""
{total}
customer messages
{ready}
ready orders
{followups}
follow-ups needed
""" def build_review_board(data): ready_cards = [] followup_cards = [] for order in data.get("orders", []): customer = html.escape(text_value(order.get("customer")) or "Customer") item = html.escape(text_value(order.get("item")) or "order") quantity = html.escape(text_value(order.get("quantity")) or "quantity to confirm") flavor = html.escape(text_value(order.get("flavor"))) pickup_time = html.escape(text_value(order.get("pickup_time"))) delivery_address = html.escape(text_value(order.get("delivery_address"))) payment_status = html.escape(text_value(order.get("payment_status"))) fields = missing_list(order) title = f"{quantity} {item}" if flavor: title += f" ({flavor})" details = [] if pickup_time: details.append(f"Time: {pickup_time}") if delivery_address: details.append(f"Place: {delivery_address}") if payment_status: details.append(f"Payment: {payment_status}") detail_html = " · ".join(details) if details else "Details still need review." if fields: chips = "".join( f'{html.escape(field.replace("_", " "))}' for field in fields ) needed = ", ".join(field.replace("_", " ") for field in fields) suggested_reply = html.escape( f"Thanks, {customer}! I have your {item} order. Could you confirm the {needed}?" ) followup_cards.append( f"""
{customer}
{title}
{detail_html}
{chips}
Suggested reply
{suggested_reply}
""" ) else: ready_cards.append( f"""
{customer}
{title}
{detail_html}
""" ) ready_html = "".join(ready_cards) or '
No ready orders yet.
' followup_html = "".join(followup_cards) or '
No follow-ups needed.
' return f"""

Ready to prep

{ready_html}

Needs follow-up

{followup_html}
""" def split_customer_messages(messages): entries = [] current_customer = "" current_parts = [] for raw_line in messages.splitlines(): line = raw_line.strip() if not line: continue if ":" in line: possible_name, body = line.split(":", 1) if possible_name.strip() and len(possible_name.strip().split()) <= 3: if current_customer or current_parts: entries.append((current_customer or "Customer", " ".join(current_parts).strip())) current_customer = possible_name.strip() current_parts = [body.strip()] continue if current_parts: current_parts.append(line) else: entries.append(("Customer", line)) if current_customer or current_parts: entries.append((current_customer or "Customer", " ".join(current_parts).strip())) return [(name, body) for name, body in entries if body] def extract_single_order(customer, message): prompt = tokenizer.apply_chat_template( [ {"role": "system", "content": SINGLE_ORDER_PROMPT}, {"role": "user", "content": f"Customer: {customer}\nMessage: {message}"}, ], tokenize=False, add_generation_prompt=True, ) inputs = tokenizer(prompt, return_tensors="pt") with torch.no_grad(): output = model.generate( **inputs, max_new_tokens=350, do_sample=False, pad_token_id=tokenizer.eos_token_id, ) generated = tokenizer.decode( output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True, ) try: parsed = extract_json(generated) except Exception: parsed = { "item": "", "quantity": "", "flavor": "", "pickup_time": "", "delivery_address": "", "payment_status": "", "notes": message, "missing_fields": [], } order = {"customer": customer} for col in ORDER_COLUMNS[1:]: value = parsed.get(col, "") if col == "missing_fields": if isinstance(value, list): order[col] = value elif isinstance(value, str): order[col] = [part.strip() for part in value.split(",") if part.strip()] else: order[col] = [] else: order[col] = text_value(value) return post_process_order(order, message) def analyze_messages(messages): if not messages.strip(): return pd.DataFrame(columns=ORDER_COLUMNS), "Choose a scenario and organize the inbox.", "", "", "", "" entries = split_customer_messages(messages) orders_data = [extract_single_order(customer, message) for customer, message in entries] data = {"orders": orders_data} orders_df = normalize_orders(data) auto_prep = build_prep_list(data) auto_replies = build_reply_drafts(data) data["prep_list"] = auto_prep data["reply_drafts"] = auto_replies summary = build_summary(data) review = build_review_board(data) prep = format_list("Prep list", auto_prep) replies = format_replies(auto_replies) raw = json.dumps(data, indent=2, ensure_ascii=False) return orders_df, summary, review, prep, replies, raw with gr.Blocks(title="DM Order Desk", theme=gr.themes.Soft()) as demo: gr.Markdown( f""" # DM Order Desk A small-model inbox assistant for tiny sellers. **Unofficial Build Small Hackathon-inspired project** · `{MODEL_ID}` · Gradio Space """ ) with gr.Row(): with gr.Column(scale=1): gr.Markdown("### Customer inbox") messages = gr.Textbox( label="Messy customer DMs", value=EXAMPLE_INPUT, lines=14, ) run = gr.Button("Organize orders", variant="primary") gr.Markdown("### Demo scenarios") with gr.Row(): home_btn = gr.Button("Home Bakery") food_btn = gr.Button("Food Truck") with gr.Row(): craft_btn = gr.Button("Craft Market") farm_btn = gr.Button("Farmers Market") with gr.Column(scale=2): summary = gr.HTML("Choose a scenario and organize the inbox.") with gr.Tabs(): with gr.Tab("Review Board"): review = gr.HTML() with gr.Tab("Data Table"): orders = gr.Dataframe(label="Orders", headers=ORDER_COLUMNS) with gr.Tab("Prep List"): prep = gr.Markdown() with gr.Tab("Reply Drafts"): replies = gr.Markdown() with gr.Tab("Raw JSON"): raw = gr.Code(label="Raw JSON", language="json") run.click(analyze_messages, inputs=messages, outputs=[orders, summary, review, prep, replies, raw]) home_btn.click(lambda: EXAMPLE_INPUT, outputs=messages) food_btn.click(lambda: FOOD_TRUCK_EXAMPLE, outputs=messages) craft_btn.click(lambda: CRAFT_SELLER_EXAMPLE, outputs=messages) farm_btn.click(lambda: FARMERS_MARKET_EXAMPLE, outputs=messages) demo.launch()