Update app.py

app.py

@@ -2,82 +2,149 @@ import gradio as gr
 import json
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from peft import PeftModel
 
 # ── Load model once at startup ──────────────────────────────
 BASE_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
-LORA_MODEL = "suneeldk/json-extract"  # ← change this
+LORA_MODEL = "suneeldk/json-extract" 
 
 tokenizer = AutoTokenizer.from_pretrained(LORA_MODEL)
 
+# Load in 4-bit for faster inference
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16,
+)
+
 base_model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
-    torch_dtype=torch.float16,
+    quantization_config=bnb_config,
     device_map="auto",
 )
+
 model = PeftModel.from_pretrained(base_model, LORA_MODEL)
+model = model.merge_and_unload()  # Merge LoRA into base — removes adapter overhead
 model.eval()
 
 
+# ── Auto-detect schema from text ────────────────────────────
+def auto_schema(text):
+    text_lower = text.lower()
+    schema = {}
+
+    money_keywords = ["paid", "sent", "received", "cost", "price", "rupees", "rs",
+                      "₹", "$", "bought", "sold", "charged", "fee", "salary",
+                      "budget", "owes", "owe", "lent", "borrowed", "fare", "rent"]
+    if any(k in text_lower for k in money_keywords) or any(c.isdigit() for c in text):
+        schema["amount"] = "number|null"
+
+    person_keywords = ["to", "from", "with", "for", "by", "told", "asked",
+                       "met", "called", "emailed", "messaged", "owes", "owe"]
+    if any(k in text_lower for k in person_keywords):
+        schema["person"] = "string|null"
+
+    date_keywords = ["jan", "feb", "mar", "apr", "may", "jun", "jul", "aug",
+                     "sep", "oct", "nov", "dec", "monday", "tuesday", "wednesday",
+                     "thursday", "friday", "saturday", "sunday", "today", "tomorrow",
+                     "yesterday", "morning", "evening", "night", "on", "at", "pm", "am"]
+    if any(k in text_lower for k in date_keywords):
+        schema["date"] = "ISO date|null"
+        if any(k in text_lower for k in ["pm", "am", "morning", "evening", "night", "at"]):
+            schema["time"] = "string|null"
+
+    item_keywords = ["bought", "ordered", "purchased", "delivered", "shipped",
+                     "kg", "litre", "pieces", "items", "pack", "bottle"]
+    if any(k in text_lower for k in item_keywords):
+        schema["item"] = "string|null"
+        schema["quantity"] = "string|null"
+
+    location_keywords = ["from", "to", "at", "in", "store", "shop", "restaurant",
+                         "station", "airport", "hotel", "office", "train", "flight", "bus"]
+    if any(k in text_lower for k in location_keywords):
+        schema["location"] = "string|null"
+
+    travel_keywords = ["train", "flight", "bus", "booked", "ticket", "pnr",
+                       "travel", "trip", "journey"]
+    if any(k in text_lower for k in travel_keywords):
+        schema["from_location"] = "string|null"
+        schema["to_location"] = "string|null"
+        schema.pop("location", None)
+
+    meeting_keywords = ["meeting", "call", "discuss", "review", "presentation",
+                        "interview", "appointment", "schedule"]
+    if any(k in text_lower for k in meeting_keywords):
+        schema["topic"] = "string|null"
+
+    schema["note"] = "string|null"
+
+    if len(schema) <= 1:
+        schema = {
+            "amount": "number|null",
+            "person": "string|null",
+            "date": "ISO date|null",
+            "note": "string|null",
+        }
+
+    return schema
+
+
 # ── Inference function ──────────────────────────────────────
 @spaces.GPU
-def extract(text, schema_text):
+def extract(text, custom_schema):
     if not text.strip():
-        return "Please enter some text."
-    if not schema_text.strip():
-        return "Please enter a schema."
-
-    try:
-        schema = json.loads(schema_text)
-    except json.JSONDecodeError:
-        return "Invalid JSON schema. Please check the format."
-
-    prompt = f"### Input: {text}\n### Schema: {json.dumps(schema)}\n### Output:"
+        return "", ""
+
+    if custom_schema and custom_schema.strip():
+        try:
+            schema = json.loads(custom_schema)
+        except json.JSONDecodeError:
+            return "Invalid JSON schema.", ""
+    else:
+        schema = auto_schema(text)
+
+    schema_str = json.dumps(schema)
+    prompt = f"### Input: {text}\n### Schema: {schema_str}\n### Output:"
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=512,
-            temperature=0.1,
-            do_sample=True,
+            max_new_tokens=128,       # JSON output is short, no need for 512
+            do_sample=False,           # Greedy decoding — faster than sampling
             pad_token_id=tokenizer.eos_token_id,
         )
 
-    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    output_part = result.split("### Output:")[-1].strip()
+    # Decode only the new tokens, skip the prompt
+    new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
+    output_part = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
 
     try:
         parsed = json.loads(output_part)
-        return json.dumps(parsed, indent=2, ensure_ascii=False)
+        return json.dumps(parsed, indent=2, ensure_ascii=False), json.dumps(schema, indent=2)
     except json.JSONDecodeError:
-        return output_part
+        return output_part, json.dumps(schema, indent=2)
 
 
 # ── Example inputs ──────────────────────────────────────────
 examples = [
-    [
-        "Paid 500 to Ravi for lunch on Jan 5",
-        '{"amount": "number", "person": "string|null", "date": "ISO date|null", "note": "string|null"}',
-    ],
-    [
-        "Meeting with Sarah at 3pm tomorrow to discuss the project budget of $10,000",
-        '{"person": "string", "time": "string", "topic": "string", "budget": "number|null"}',
-    ],
-    [
-        "Bought 3 kg of rice from Krishna Stores for 250 rupees on March 10",
-        '{"item": "string", "quantity": "string", "store": "string", "amount": "number", "date": "ISO date|null"}',
-    ],
+    ["Paid 500 to Ravi for lunch on Jan 5"],
+    ["Meeting with Sarah at 3pm tomorrow to discuss the project budget of $10,000"],
+    ["Bought 3 kg of rice from Krishna Stores for 250 rupees on March 10"],
+    ["Booked a train from Chennai to Bangalore on April 10 for 750 rupees"],
+    ["Ravi owes me 300 for last week's dinner"],
+    ["Ordered 2 pizzas and 1 coke from Dominos for 850 rupees"],
 ]
 
 # ── Gradio UI ───────────────────────────────────────────────
-with gr.Blocks(title="json-extract", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="json-extract") as demo:
     gr.Markdown(
         """
         # json-extract
         Extract structured JSON from natural language text.
-        Enter any text and a target JSON schema — the model returns clean JSON output.
+
+        Just type a sentence — the model auto-detects the right schema and extracts clean JSON.
         """
     )
 
@@ -88,22 +155,25 @@ with gr.Blocks(title="json-extract", theme=gr.themes.Soft()) as demo:
                 placeholder="e.g. Paid 500 to Ravi for lunch on Jan 5",
                 lines=3,
             )
-            schema_input = gr.Textbox(
-                label="JSON Schema",
-                placeholder='e.g. {"amount": "number", "person": "string|null"}',
-                lines=3,
-            )
             btn = gr.Button("Extract", variant="primary")
 
+            with gr.Accordion("Advanced: Custom Schema (optional)", open=False):
+                schema_input = gr.Textbox(
+                    label="Custom JSON Schema",
+                    placeholder='Leave empty for auto-detect, or enter e.g. {"amount": "number", "person": "string|null"}',
+                    lines=3,
+                )
+
         with gr.Column():
             output = gr.Textbox(label="Extracted JSON", lines=10)
+            detected_schema = gr.Textbox(label="Schema Used", lines=5)
 
     gr.Examples(
         examples=examples,
-        inputs=[text_input, schema_input],
+        inputs=[text_input],
     )
 
-    btn.click(fn=extract, inputs=[text_input, schema_input], outputs=output)
-    text_input.submit(fn=extract, inputs=[text_input, schema_input], outputs=output)
+    btn.click(fn=extract, inputs=[text_input, schema_input], outputs=[output, detected_schema])
+    text_input.submit(fn=extract, inputs=[text_input, schema_input], outputs=[output, detected_schema])
 
 demo.launch()