Handle device and dtype better

Makefile

@@ -10,7 +10,7 @@ clean:
 	rm -rf .venv
 
 run:
-	source .venv/bin/activate ; python address_parser.py
+	source .venv/bin/activate ; python app.py
 
 black-check:
 	source .venv/bin/activate ; black . --check --verbose --line-length 120

app.py

@@ -3,16 +3,36 @@ from transformers import pipeline
 import gradio as gr
 import json
 
-# Initialize model pipeline
 model_id = "dylanhogg/gnaf-structured-address-v0.1-75a1791-20250921-063650"
+max_new_tokens = 256
+do_sample = False
+
+if torch.cuda.is_available():
+    device = "cuda"
+    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+    device_map = "auto"
+    print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+elif torch.backends.mps.is_available():
+    device = "mps"
+    dtype = torch.float16
+    device_map = None
+else:
+    device = "cpu"
+    dtype = torch.float32
+    device_map = None
+
+print(f"Device settings: {device=}, {dtype=}, {device_map=}")
 
 pipe = pipeline(
     "text-generation",
     model=model_id,
-    dtype=torch.bfloat16,
-    device_map="auto",
+    dtype=dtype,
+    device_map=device_map,
 )
 
+print(f"Model {pipe.model=}")
+print(f"Model config {pipe.model.config=}")
+
 
 def parse_address(user_address: str) -> tuple[str, str]:
     """Parse address and return both raw response and JSON"""
@@ -24,8 +44,8 @@ def parse_address(user_address: str) -> tuple[str, str]:
 
     outputs = pipe(
         messages,
-        max_new_tokens=256,
-        do_sample=False,
+        max_new_tokens=max_new_tokens,
+        do_sample=do_sample,
     )
 
     response = outputs[0]["generated_text"]
@@ -48,7 +68,6 @@ def parse_address(user_address: str) -> tuple[str, str]:
     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.")
@@ -64,7 +83,6 @@ with gr.Blocks(title="Address Parser") as demo:
             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",
@@ -99,10 +117,10 @@ with gr.Blocks(title="Address Parser") as demo:
         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__":
+    print("Launching app...")
     demo.launch()
+    print("Done.")