Update app.py

app.py

@@ -1,4 +1,5 @@
 import re
+import json
 import torch
 import gradio as gr
 from transformers import T5Tokenizer, T5ForConditionalGeneration
@@ -17,6 +18,13 @@ def extract_amount(input_text):
     match = re.search(r'\b\d{1,15}([.,]\d{1,15})?\b', input_text)
     return match.group(0) if match else "Not found"  # Keep as string
 
+def parse_model_output(model_output):
+    """Attempts to parse the model's output into a dictionary."""
+    try:
+        return json.loads(model_output)  # Safely parse JSON
+    except json.JSONDecodeError:
+        return {}  # Return an empty dictionary if parsing fails
+
 def generate_command(input_command):
     """Runs the model and extracts information."""
     input_ids = tokenizer("extract: " + input_command, return_tensors="pt").input_ids.to(device)
@@ -27,13 +35,13 @@ def generate_command(input_command):
     # Extract amount manually
     extracted_amount = extract_amount(input_command)
 
-    try:
-        model_data = eval(model_output)  # Convert model output into a Python dictionary
-        action = model_data.get("action", "Not found")
-        currency = model_data.get("currency", "Not found")
-        recipient = model_data.get("recipient", "Not found")
-    except:
-        action, currency, recipient = "Error", "Error", "Error"
+    # Parse the model output safely
+    model_data = parse_model_output(model_output)
+
+    # Extract required variables
+    action = model_data.get("action", "Not found")
+    currency = model_data.get("currency", "Not found")
+    recipient = model_data.get("recipient", "Not found")
 
     return extracted_amount, action, currency, recipient