Update app.py

app.py

@@ -2,24 +2,31 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import json
-import time
 
-# Model Loading
-tokenizer = AutoTokenizer.from_pretrained("numind/NuExtract-1.5")
-model = AutoModelForCausalLM.from_pretrained(
-    "numind/NuExtract-1.5",
-    device_map="auto",
-    torch_dtype=torch.float16
-)
+# Initialize with error handling
+try:
+    tokenizer = AutoTokenizer.from_pretrained("numind/NuExtract-1.5")
+    model = AutoModelForCausalLM.from_pretrained(
+        "numind/NuExtract-1.5",
+        device_map="auto",
+        torch_dtype=torch.float16
+    )
+    MODEL_LOADED = True
+except Exception as e:
+    MODEL_LOADED = False
+    print(f"Model loading failed: {e}")
 
 def extract_structure(template, text):
-    prompt = f"""Extract the following fields from the text:
+    if not MODEL_LOADED:
+        return "❌ Model not loaded", {}, "<p style='color:red'>Model failed to initialize</p>"
+    
+    prompt = f"""Extract from text:
 Template: {template}
 Text: {text}
-Extracted JSON:"""
+JSON Output:"""
     
     try:
-        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
         outputs = model.generate(**inputs, max_new_tokens=512)
         result = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
@@ -33,21 +40,10 @@ Extracted JSON:"""
     except Exception as e:
         return f"❌ Error: {str(e)}", {}, f"<p style='color:red'>{str(e)}</p>"
 
-# Gradio Interface
+# Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# NuExtract-1.5 Structured Data Extractor")
-    
-    with gr.Row():
-        with gr.Column():
-            template = gr.Textbox(label="Template (JSON)", value='{"fields": ["name", "email"]}')
-            text = gr.TextArea(label="Input Text")
-            btn = gr.Button("Extract")
-        
-        with gr.Column():
-            status = gr.Textbox(label="Status")
-            json_out = gr.JSON(label="Output")
-            html_out = gr.HTML()
-    
-    btn.click(extract_structure, [template, text], [status, json_out, html_out])
+    # [Keep your existing UI code here]
+    # ...
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()