dm-order-desk / app.py
SSSSSSSiao's picture
Update app.py
852582b verified
Raw
History Blame
19.3 kB
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"""
<div style="display:grid;grid-template-columns:repeat(3,minmax(0,1fr));gap:12px;margin:8px 0 18px;">
<div style="background:#f8fafc;border:1px solid #e2e8f0;border-radius:10px;padding:14px;">
<div style="font-size:24px;font-weight:700;">{total}</div>
<div style="color:#475569;">customer messages</div>
</div>
<div style="background:#ecfdf5;border:1px solid #bbf7d0;border-radius:10px;padding:14px;">
<div style="font-size:24px;font-weight:700;">{ready}</div>
<div style="color:#166534;">ready orders</div>
</div>
<div style="background:#fffbeb;border:1px solid #fde68a;border-radius:10px;padding:14px;">
<div style="font-size:24px;font-weight:700;">{followups}</div>
<div style="color:#92400e;">follow-ups needed</div>
</div>
</div>
"""
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'<span style="display:inline-block;background:#fef3c7;color:#92400e;border-radius:999px;padding:4px 9px;margin:3px;font-size:12px;">{html.escape(field.replace("_", " "))}</span>'
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"""
<div style="background:#fff7ed;border:1px solid #fed7aa;border-radius:10px;padding:14px;margin-bottom:10px;">
<div style="font-weight:700;">{customer}</div>
<div>{title}</div>
<div style="color:#64748b;font-size:13px;margin-top:4px;">{detail_html}</div>
<div style="margin-top:8px;">{chips}</div>
<div style="margin-top:10px;background:#ffffff;border:1px solid #fed7aa;border-radius:8px;padding:10px;color:#334155;font-size:13px;">
<div style="font-weight:700;margin-bottom:4px;">Suggested reply</div>
{suggested_reply}
</div>
</div>
"""
)
else:
ready_cards.append(
f"""
<div style="background:#f0fdf4;border:1px solid #bbf7d0;border-radius:10px;padding:14px;margin-bottom:10px;">
<div style="font-weight:700;">{customer}</div>
<div>{title}</div>
<div style="color:#64748b;font-size:13px;margin-top:4px;">{detail_html}</div>
</div>
"""
)
ready_html = "".join(ready_cards) or '<div style="color:#64748b;">No ready orders yet.</div>'
followup_html = "".join(followup_cards) or '<div style="color:#64748b;">No follow-ups needed.</div>'
return f"""
<div style="display:grid;grid-template-columns:1fr 1fr;gap:16px;">
<section>
<h3>Ready to prep</h3>
{ready_html}
</section>
<section>
<h3>Needs follow-up</h3>
{followup_html}
</section>
</div>
"""
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()