Update app.py

app.py

@@ -43,4 +43,57 @@ def extract_data(raw_text, fields_to_extract):
         # Fallback: Safely strip markdown code blocks without using complex regex
         # that might break code editors during copy-pasting
         if output_text.startswith("```"):
-            # Remove the starting ```json or
+            # Remove the starting ```json or
+            # Remove the ending ```
+            output_text = re.sub(r"\n?```$", "", output_text)
+
+        # Parse the text into an actual JSON dictionary for the Gradio UI
+        structured_data = json.loads(output_text.strip())
+        return structured_data
+
+    except json.JSONDecodeError:
+        return {
+            "error": "The model failed to return valid JSON. It returned this instead:",
+            "raw_output": output_text
+        }
+    except Exception as e:
+        return {"error": str(e)}
+
+# -------------------------
+# Build the Gradio UI
+# -------------------------
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🛟 The Data Rescuer")
+    gr.Markdown(f"**Powered by `{MODEL_ID}`** | Turn messy transcripts, notes, and OCR text into clean JSON data.")
+    
+    with gr.Row():
+        # Left Column: Inputs
+        with gr.Column():
+            raw_input = gr.Textbox(
+                label="1. Paste Unstructured Text",
+                placeholder="Paste your messy meeting notes, emails, or raw text here...",
+                lines=12
+            )
+            
+            schema_input = gr.Textbox(
+                label="2. What fields do you want to extract?",
+                placeholder="e.g., Company Name, Contact Person, Deadline, Action Items (list)",
+                lines=3
+            )
+            
+            extract_btn = gr.Button("Extract Structured Data", variant="primary")
+            
+        # Right Column: Output
+        with gr.Column():
+            json_output = gr.JSON(label="Structured Output")
+
+    # Connect the button to the function
+    extract_btn.click(
+        fn=extract_data,
+        inputs=[raw_input, schema_input],
+        outputs=json_output
+    )
+
+# Launch the app
+if __name__ == "__main__":
+    demo.launch()