Update app.py

app.py

@@ -1,6 +1,6 @@
 import json
-import torch
 import re
+import torch
 import gradio as gr
 from transformers import T5Tokenizer, T5ForConditionalGeneration
 
@@ -9,74 +9,65 @@ model_name = "./t5-finetuned-final"
 tokenizer = T5Tokenizer.from_pretrained(model_name)
 model = T5ForConditionalGeneration.from_pretrained(model_name)
 
-# Move model to GPU if available
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-# Enable optimizations for GPU
 if torch.cuda.is_available():
-    model.half()  # Use half-precision for faster computation
+    model.half()
     try:
-        model = torch.compile(model)  # PyTorch 2.0+ optimization
-    except Exception:
-        pass  # Ignore if torch.compile is not available
+        model = torch.compile(model)
+    except:
+        pass
 
 def correct_amount_format(output):
-    """
-    This function attempts to fix the numeric formatting issues in the generated output:
-      1. It replaces a comma used as a decimal separator (i.e. followed by exactly two digits) with a period.
-      2. It converts the number to a float and rounds it to two decimal places.
-      
-    If the output is valid JSON, it will update the "amount" field accordingly.
-    Otherwise, it falls back to a regex-based fix.
-    """
+    # Attempt to parse as JSON and correct amounts
     try:
-        # Try to parse the output as JSON
         data = json.loads(output)
-        if "amount" in data and isinstance(data["amount"], str):
-            # Replace a comma that is likely a decimal separator (e.g., "10,50" -> "10.50")
-            amount_str = re.sub(r'(\d+),(\d{2})\b', r'\1.\2', data["amount"])
-            try:
-                # Convert to float, round to two decimals, then reformat
-                num = float(amount_str)
-                rounded = round(num, 2)
-                data["amount"] = "{:.2f}".format(rounded)
-            except ValueError:
-                # If conversion fails, leave the original value
-                pass
+        
+        def correct_value(value):
+            if isinstance(value, str):
+                # Remove commas used as thousand separators
+                value = re.sub(r',(?=\d{3})', '', value)
+                # Replace the first comma with a period (decimal)
+                value = value.replace(',', '.', 1)
+            return value
+        
+        # Correct each value in the JSON data
+        if isinstance(data, dict):
+            for key in data:
+                data[key] = correct_value(data[key])
+        elif isinstance(data, list):
+            for i in range(len(data)):
+                data[i] = correct_value(data[i])
+        
         return json.dumps(data, ensure_ascii=False)
     except json.JSONDecodeError:
-        # Fallback if output is not valid JSON:
-        # Replace commas used as decimal separators (only if followed by exactly 2 digits)
-        output = re.sub(r'(\d+),(\d{2})\b', r'\1.\2', output)
-        # Fallback: truncate any extra digits (note: this does not round)
-        output = re.sub(r'(\d+\.\d{2})\d+', r'\1', output)
+        # Fallback for invalid JSON: basic corrections
+        output = re.sub(r'(\d+),(\d+)\b', r'\1.\2', output)
         return output
 
 def generate_command(input_command):
     prompt = "extract: " + input_command
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
-
+    
     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,
+        num_beams=3,
         early_stopping=True
     )
-
+    
     result = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    # Apply the updated post-processing to correct the amount formatting
     result = correct_amount_format(result)
     return result
 
-# Create a Gradio interface
 iface = gr.Interface(
     fn=generate_command,
     inputs=gr.Textbox(lines=2, placeholder="Enter a command..."),
     outputs=gr.Textbox(label="Extracted JSON Output"),
     title="T5 Fine-Tuned Command Extractor",
-    description="Enter a command, and the fine-tuned T5 model will extract relevant details in JSON format.",
+    description="Extracts details in JSON format with exact amount preservation.",
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()