Improve extraction prompt and model

app.py

@@ -4,7 +4,8 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
+# MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
+MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
 
 ORDER_COLUMNS = [
     "customer",
@@ -22,9 +23,41 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32)
 model.eval()
 
+# SYSTEM_PROMPT = """
+# You extract customer orders from messy DMs for tiny sellers.
+# 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": [
+#     {"customer": "", "reply": ""}
+#   ]
+# }
+# Use empty strings for unknown values. Put missing details in missing_fields.
+# """
+
+# 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.
+# """
+
 SYSTEM_PROMPT = """
-You extract customer orders from messy DMs for tiny sellers.
-Return only valid JSON with this exact shape:
+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": [
     {
@@ -44,13 +77,19 @@ Return only valid JSON with this exact shape:
     {"customer": "", "reply": ""}
   ]
 }
-Use empty strings for unknown values. Put missing details in missing_fields.
-"""
 
-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.
+Critical rules:
+- Do not invent customer names.
+- Do not change customer names.
+- Do not merge different customers.
+- Include every customer message that looks like an order or possible order.
+- Use only facts explicitly present in the messages.
+- If a value is unknown, use an empty string.
+- Do not add order_id or total_cost unless the message mentions them.
+- For pickup orders, put the pickup time in pickup_time and leave delivery_address empty unless a real address is given.
+- If the customer is unsure, still include the order and put uncertainty in notes.
+- missing_fields should only include practical fields the seller needs to ask for, such as quantity, flavor, pickup_time, delivery_address, or payment_status.
+- Reply drafts should be short, friendly, and ask only for missing information.
 """
 
 def extract_json(text):