File size: 2,907 Bytes
44c8200 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
import torch
from transformers import pipeline
import gradio as gr
import json
# Initialize model pipeline
model_id = "dylanhogg/gnaf-structured-address-v0.1-75a1791-20250921-063650"
pipe = pipeline(
"text-generation",
model=model_id,
dtype=torch.bfloat16,
device_map="auto",
)
def parse_address(user_address: str) -> tuple[str, str]:
"""Parse address and return both raw response and JSON"""
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},
]
outputs = pipe(
messages,
max_new_tokens=256,
do_sample=False,
)
response = outputs[0]["generated_text"]
last_content = response[-1]["content"]
# Try to extract JSON from the response
try:
# Find first '{' and last '}' to extract JSON
start = last_content.find("{")
end = last_content.rfind("}") + 1
if start != -1 and end > start:
json_str = last_content[start:end]
parsed = json.loads(json_str)
formatted_json = json.dumps(parsed, indent=2)
else:
formatted_json = "No JSON found in response"
except Exception as e:
formatted_json = f"Error parsing JSON: {str(e)}\n\nRaw output:\n{last_content}"
return formatted_json, last_content
# Create Gradio interface
with gr.Blocks(title="Address Parser") as demo:
gr.Markdown("# 🏠 Structured Address Parser")
gr.Markdown("This model converts text addresses into structured JSON format.")
with gr.Row():
with gr.Column():
input_text = gr.Textbox(
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)
# Examples
gr.Examples(
examples=[
"48a Pirrama Rd Pyrmont NSW 2009",
"Floor 3, 152-156 Clarence St, Sydney NSW 2000",
"Aptt 16, 400 Bondi Rd, Bondi NSW 2026",
"Unit 18/14-18 Flood St, Bondi, NSW 2026",
"Lvl 15/333 George St Sydney NSW 2000",
"Check out: 44 Ulm St, Maroubra NSW 2035",
"44 Ulm St, Maroubra NSW 2035 is where it's at!",
],
inputs=input_text,
)
# Handle events
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])
# Launch the app
if __name__ == "__main__":
demo.launch()
|