Update app.py

app.py

@@ -1,5 +1,6 @@
 import json
 import torch
+import re
 import gradio as gr
 from transformers import T5Tokenizer, T5ForConditionalGeneration
 
@@ -20,22 +21,33 @@ if torch.cuda.is_available():
     except:
         pass  # Ignore if torch.compile is not available
 
-# Define the function for inference
+# Corrects formatting of amounts, handles commas and precision issues
+def correct_amount_format(output):
+    # Fix the commas in decimal amounts
+    output = re.sub(r'(\d+),(\d+)', r'\1.\2', output)
+    
+    # Ensure that numbers with more than 2 decimal places are rounded
+    output = re.sub(r'(\d+\.\d{2})\d+', r'\1', output)  # Keeps only 2 decimal places
+    return output
+
+# Function to generate command and parse amounts
 def generate_command(input_command):
     prompt = "extract: " + input_command
-
-    # Encode input and move tensors to the correct device
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
 
-    # Generate with optimized settings
     output_ids = model.generate(
         input_ids,
         max_length=64,  # Reduced for speed
-        num_beams=3,  # Lowered from 5 to 3 for faster output
+        num_beams=3,    # Lowered from 5 to 3 for faster output
         early_stopping=True
     )
 
-    return tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    result = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+    # Apply the post-processing fix to the amount
+    result = correct_amount_format(result)
+
+    return result
 
 # Create a Gradio interface
 iface = gr.Interface(