Update app.py

app.py

@@ -1,3 +1,4 @@
+import re
 import torch
 import gradio as gr
 from transformers import T5Tokenizer, T5ForConditionalGeneration
@@ -11,23 +12,44 @@ model = T5ForConditionalGeneration.from_pretrained(model_name)
 device = torch.device("cpu")
 model.to(device)
 
+def extract_amount(input_text):
+    """Finds the first number in the text (supports both . and , as decimal separators)."""
+    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 generate_command(input_command):
-    """Runs the model and returns its output as plain text."""
+    """Runs the model and extracts information."""
     input_ids = tokenizer("extract: " + input_command, return_tensors="pt").input_ids.to(device)
-
-    output_ids = model.generate(input_ids, max_length=128, num_beams=1, early_stopping=True)  # Simple decoding
+    
+    output_ids = model.generate(input_ids, max_length=128, num_beams=1, early_stopping=True)
     model_output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
-    return model_output  # Return the model's raw output
+    # 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"
 
+    return extracted_amount, action, currency, recipient
+
+# Gradio Interface
 iface = gr.Interface(
     fn=generate_command,
     inputs=gr.Textbox(lines=2, placeholder="Enter a command..."),
-    outputs=gr.Textbox(label="Model Output"),
-    title="T5 Model Output Test",
-    description="This version only shows the model's raw output."
+    outputs=[
+        gr.Textbox(label="Extracted Amount"),  
+        gr.Textbox(label="Action"),
+        gr.Textbox(label="Currency"),
+        gr.Textbox(label="Recipient"),
+    ],
+    title="Command Extractor",
+    description="Extracts amount separately and displays action, currency, and recipient from model output."
 )
 
 if __name__ == "__main__":
     iface.launch()
-