Update app.py

app.py

@@ -1,71 +1,74 @@
-import re
+import json
 import torch
+import re
 import gradio as gr
 from transformers import T5Tokenizer, T5ForConditionalGeneration
 
-# Load the fine-tuned model and tokenizer from the local folder
+# Load the fine-tuned model
 model_name = "./t5-finetuned-final"
 tokenizer = T5Tokenizer.from_pretrained(model_name)
 model = T5ForConditionalGeneration.from_pretrained(model_name)
 
-# Move model to GPU if available; otherwise, it will run on CPU
+# Move model to GPU if available
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-# If using GPU, enable half precision and try torch.compile (PyTorch 2.0+)
+# Enable optimizations for GPU
 if torch.cuda.is_available():
-    model.half()  # Use half-precision for faster computation on GPU
+    model.half()  # Use half-precision for faster computation
     try:
-        model = torch.compile(model)
-    except Exception:
-        pass  # Continue if torch.compile is unavailable
+        model = torch.compile(model)  # PyTorch 2.0+ optimization
+    except:
+        pass  # Ignore if torch.compile is not available
 
-def fix_amount_in_output(input_command, output_str):
+def sanitize_amount(output):
     """
-    Extracts the first decimal number from the input and replaces the "amount" value
-    in the output with that exact value.
+    Sanitizes the amount field to ensure it is correctly formatted.
     """
-    # Look for a number with optional decimal separator in the input command.
-    match = re.search(r'(\d+(?:[.,]\d+))', input_command)
-    if match:
-        # Normalize any commas to a period.
-        correct_amount_str = match.group(1).replace(',', '.')
-    else:
-        return output_str
+    # Fix malformed amounts like "46307.0" -> "4630.07" or "4630327.0" -> "463032.07"
+    def fix_malformed_amount(match):
+        full_match = match.group(0)
+        integer_part = match.group(1)
+        decimal_part = match.group(2)
+        return f"{integer_part}.{decimal_part}"  # Reconstruct the correct format
 
-    # Replace the "amount" value in the output with the extracted amount.
-    fixed_output = re.sub(
-        r'("amount"\s*:\s*)(\d+(?:\.\d+)?)',
-        r'\1' + correct_amount_str,
-        output_str
-    )
-    return fixed_output
+    # Match numbers with misplaced decimal points
+    output = re.sub(r'(\d+)(\d{2})\.0', fix_malformed_amount, output)
+    return output
 
 def generate_command(input_command):
+    """
+    Generates the command and ensures the exact amount is displayed without changes.
+    """
     prompt = "extract: " + input_command
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
-
-    # Use greedy decoding (num_beams=1) on CPU for speed; otherwise, use beam search on GPU.
-    if device.type == "cpu":
-        output_ids = model.generate(
-            input_ids,
-            max_length=64,
-            num_beams=1,         # Greedy decoding for faster output on CPU
-            early_stopping=True
-        )
-    else:
-        output_ids = model.generate(
-            input_ids,
-            max_length=64,
-            num_beams=3,         # Beam search for potentially higher quality on GPU
-            early_stopping=True
-        )
-
+    
+    # Generate output from the model
+    output_ids = model.generate(
+        input_ids,
+        max_length=64,  # Reduced for speed
+        num_beams=3,    # Lowered from 5 to 3 for faster output
+        early_stopping=True
+    )
     result = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    result_fixed = fix_amount_in_output(input_command, result)
-    return result_fixed
+    
+    # Sanitize the output to fix malformed amounts
+    sanitized_result = sanitize_amount(result)
+    
+    try:
+        # Attempt to parse the sanitized result as JSON
+        data = json.loads(sanitized_result)
+        
+        # Convert numeric amounts to strings to preserve exact formatting
+        if isinstance(data.get("amount"), (int, float)):
+            data["amount"] = str(data["amount"])
+        
+        return json.dumps(data, ensure_ascii=False)  # Return as JSON string
+    except json.JSONDecodeError:
+        # If not valid JSON, return the raw sanitized output
+        return sanitized_result
 
-# Create a Gradio interface.
+# Create a Gradio interface
 iface = gr.Interface(
     fn=generate_command,
     inputs=gr.Textbox(lines=2, placeholder="Enter a command..."),
@@ -74,8 +77,6 @@ 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()
-
-
-
+    iface.launch()