Update app.py

app.py

@@ -37,28 +37,22 @@ def fix_json_output(output):
     output = re.sub(r':\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*([,}])', r':"\1"\2', output)
     return output
 
-def replace_amount_in_json(output, amount):
+def merge_json_with_amount(model_output, amount):
     """
-    Replaces the amount in the model's JSON output with the extracted amount.
+    Merges the model's JSON output with the extracted amount.
     """
     try:
         # Fix JSON formatting issues
-        output = fix_json_output(output)
-        data = json.loads(output)
+        model_output = fix_json_output(model_output)
+        data = json.loads(model_output)
         
-        # Replace the amount field if it exists
-        if "amount" in data:
+        # Replace the amount field with the extracted amount
+        if amount:
             data["amount"] = float(amount) if '.' in amount else int(amount)
         
-        # Ensure all required fields are present
-        required_fields = ["action", "amount", "currency", "recipient"]
-        for field in required_fields:
-            if field not in data:
-                data[field] = "unknown"  # Fallback value
-        
         return json.dumps(data, ensure_ascii=False)
     except json.JSONDecodeError:
-        # Fallback for invalid JSON: return a default structure
+        # If the model's output is invalid, return a default structure with the correct amount
         return json.dumps({
             "action": "unknown",
             "amount": float(amount) if '.' in amount else int(amount),
@@ -82,11 +76,14 @@ def generate_command(input_command):
         early_stopping=True
     )
     
-    result = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    # Decode the model's output
+    model_output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
     
-    # Replace the model's amount with the extracted amount
+    # Merge the model's output with the extracted amount
     if amount:
-        result = replace_amount_in_json(result, amount)
+        result = merge_json_with_amount(model_output, amount)
+    else:
+        result = model_output  # Use the model's output as-is if no amount is found
     
     return result