Add inference time

Makefile

@@ -12,6 +12,9 @@ clean:
 run:
 	source .venv/bin/activate ; python app.py
 
+debug:
+	source .venv/bin/activate ; gradio app.py
+
 black-check:
 	source .venv/bin/activate ; black . --check --verbose --line-length 120
 

app.py

@@ -2,6 +2,7 @@ import torch
 from transformers import pipeline
 import gradio as gr
 import json
+import time
 
 model_id = "dylanhogg/gnaf-structured-address-v0.1-75a1791-20250921-063650"
 max_new_tokens = 256
@@ -32,26 +33,29 @@ pipe = pipeline(
 
 print(f"Model {pipe.model=}")
 print(f"Model config {pipe.model.config=}")
+print("Ready to parse addresses!")
 
 
-def parse_address(user_address: str) -> tuple[str, str]:
-    """Parse address and return both raw response and JSON"""
+def parse_address(user_address: str) -> tuple[str, str, str]:
+    """Parse address and return both raw response, JSON, and timing info"""
     user_content = f"Translate a text address into structured json.\n{user_address}"
     messages = [
         {"role": "system", "content": "Translate a text address into structured json."},
         {"role": "user", "content": user_content},
     ]
 
+    inference_start = time.time()
     outputs = pipe(
         messages,
         max_new_tokens=max_new_tokens,
         do_sample=do_sample,
     )
+    inference_time = time.time() - inference_start
 
     response = outputs[0]["generated_text"]
     last_content = response[-1]["content"]
 
-    # Try to extract JSON from the response
+    # Parse JSON (no timing needed)
     try:
         # Find first '{' and last '}' to extract JSON
         start = last_content.find("{")
@@ -65,23 +69,29 @@ def parse_address(user_address: str) -> tuple[str, str]:
     except Exception as e:
         formatted_json = f"Error parsing JSON: {str(e)}\n\nRaw output:\n{last_content}"
 
-    return formatted_json, last_content
+    timing_info = f"{inference_time:.3f}s (on {device} device)"
+    return formatted_json, last_content, timing_info
 
 
 with gr.Blocks(title="Address Parser") as demo:
     gr.Markdown("# 🏠 Structured Address Parser")
-    gr.Markdown("This model converts text addresses into structured JSON format.")
+    gr.Markdown("Extracts structured address information from text.")
+    gr.Markdown(
+        "Structured output is JSON format with Australian [G-NAF](https://docs.geoscape.com.au/projects/gnaf_desc/en/stable/overview.html) (Geocoded National Address File) fields."
+    )
+    gr.Markdown(f"Model: [{model_id}](https://huggingface.co/{model_id})")
 
     with gr.Row():
         with gr.Column():
             input_text = gr.Textbox(
-                label="Input Address", placeholder="Enter an address...", value="48a Pirrama Rd Pyrmont NSW 2009"
+                label="🏠 Input Address", placeholder="Enter an address...", value="48a Pirrama Rd Pyrmont NSW 2009"
             )
             submit_btn = gr.Button("Parse Address", variant="primary")
 
         with gr.Column():
-            json_output = gr.Textbox(label="Structured JSON", interactive=False, lines=10)
-            raw_output = gr.Textbox(label="Raw Model Output", interactive=False, lines=5)
+            took_output = gr.Textbox(label="⏱️ Inference Time", interactive=False, max_lines=1)
+            json_output = gr.Textbox(label="📋 Structured JSON", interactive=False, lines=10)
+            raw_output = gr.Textbox(label="🔍 Raw Model Output", interactive=False, lines=3)
 
     gr.Examples(
         examples=[
@@ -117,8 +127,8 @@ with gr.Blocks(title="Address Parser") as demo:
         inputs=input_text,
     )
 
-    submit_btn.click(fn=parse_address, inputs=input_text, outputs=[json_output, raw_output])
-    input_text.submit(fn=parse_address, inputs=input_text, outputs=[json_output, raw_output])
+    submit_btn.click(fn=parse_address, inputs=input_text, outputs=[json_output, raw_output, took_output])
+    input_text.submit(fn=parse_address, inputs=input_text, outputs=[json_output, raw_output, took_output])
 
 if __name__ == "__main__":
     print("Launching app...")