Update app.py

app.py

@@ -1,10 +1,9 @@
-import json
-import torch
 import re
+import torch
 import gradio as gr
 from transformers import T5Tokenizer, T5ForConditionalGeneration
 
-# Load the fine-tuned model
+# Load the fine-tuned model from the local folder
 model_name = "./t5-finetuned-final"
 tokenizer = T5Tokenizer.from_pretrained(model_name)
 model = T5ForConditionalGeneration.from_pretrained(model_name)
@@ -13,43 +12,59 @@ model = T5ForConditionalGeneration.from_pretrained(model_name)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-# Enable optimizations for GPU
+# Optimize for GPU: half precision and compilation (if supported)
 if torch.cuda.is_available():
     model.half()  # Use half-precision for faster computation
     try:
-        model = torch.compile(model)  # PyTorch 2.0+ optimization
-    except:
-        pass  # Ignore if torch.compile is not available
+        model = torch.compile(model)  # Optimize with torch.compile() (PyTorch 2.0+)
+    except Exception:
+        pass  # Continue if torch.compile() isn't available
 
-# 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
+def fix_amount_in_output(input_command, output_str):
+    """
+    This function extracts the first decimal number found in the input_command
+    and then replaces the "amount" field in the model output with that number.
+    """
+    # Extract the first number that has a decimal point (or comma) from the input.
+    match = re.search(r'(\d+(?:[.,]\d+))', input_command)
+    if match:
+        # Normalize to use a period as the decimal separator.
+        correct_amount_str = match.group(1).replace(',', '.')
+    else:
+        # If nothing is found, return the output unchanged.
+        return output_str
+
+    # Replace the amount value in the output.
+    # This expects the output to contain a pattern like: "amount": some_number
+    fixed_output = re.sub(
+        r'("amount"\s*:\s*)(\d+(?:\.\d+)?)',
+        r'\1' + correct_amount_str,
+        output_str
+    )
+    return fixed_output
 
-# Function to generate command and parse amounts
 def generate_command(input_command):
     prompt = "extract: " + input_command
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
 
+    # Tokenize input and send to the correct device.
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+    
+    # Generate output using optimized parameters.
     output_ids = model.generate(
         input_ids,
-        max_length=64,  # Reduced for speed
-        num_beams=3,    # Lowered from 5 to 3 for faster output
+        max_length=64,   # Reduced length for faster generation.
+        num_beams=3,     # Fewer beams for faster inference.
         early_stopping=True
     )
-
+    
+    # Decode the generated tokens.
     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
+    # Fix the "amount" field in the output using the input value.
+    result_fixed = fix_amount_in_output(input_command, result)
+    return result_fixed
 
-# Create a Gradio interface
+# Define a Gradio interface.
 iface = gr.Interface(
     fn=generate_command,
     inputs=gr.Textbox(lines=2, placeholder="Enter a command..."),
@@ -58,7 +73,7 @@ iface = gr.Interface(
     description="Enter a command, and the fine-tuned T5 model will extract relevant details in JSON format.",
 )
 
-# Launch the app
 if __name__ == "__main__":
     iface.launch()
 
+