Update app.py

app.py

@@ -19,6 +19,16 @@ if torch.cuda.is_available():
     except:
         pass
 
+def extract_amount(input_text):
+    """
+    Extracts the amount from the input text using a robust regex.
+    """
+    # Regex to match any valid number (integers or decimals)
+    amount_match = re.search(r'\d+(?:\.\d+)?', input_text)
+    if amount_match:
+        return amount_match.group(0)
+    return None
+
 def fix_json_output(output):
     """
     Fixes common JSON formatting issues in the model's output.
@@ -31,40 +41,29 @@ def fix_json_output(output):
     output = re.sub(r':\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*([,}])', r':"\1"\2', output)
     return output
 
-def correct_amount_format(output):
+def replace_amount_in_json(output, amount):
     """
-    Corrects amount formatting in the JSON output.
+    Replaces the amount in the model's JSON output with the extracted amount.
     """
     try:
         # Fix JSON formatting issues
         output = fix_json_output(output)
         data = json.loads(output)
         
-        def correct_value(value):
-            if isinstance(value, str):
-                # Fix amounts with multiple decimal points (e.g., 3140.98.0 → 3140.98)
-                if re.match(r'^\d+\.\d+\.\d+$', value):
-                    value = value.split('.')[0] + '.' + value.split('.')[1]
-                # Remove trailing .0 if it's not part of the original amount
-                if re.match(r'^\d+\.0$', value):
-                    value = value.split('.')[0]
-            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])
+        # Replace the amount field if it exists
+        if "amount" in data:
+            data["amount"] = float(amount) if '.' in amount else int(amount)
         
         return json.dumps(data, ensure_ascii=False)
     except json.JSONDecodeError:
-        # Fallback for invalid JSON: basic corrections
-        output = re.sub(r'(\d+),(\d+)\b', r'\1.\2', output)
+        # Fallback for invalid JSON: return the original output
         return output
 
 def generate_command(input_command):
+    # Extract the amount from the input
+    amount = extract_amount(input_command)
+    
+    # Generate the JSON output using the model
     prompt = "extract: " + input_command
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
     
@@ -76,7 +75,11 @@ def generate_command(input_command):
     )
     
     result = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    result = correct_amount_format(result)
+    
+    # Replace the model's amount with the extracted amount
+    if amount:
+        result = replace_amount_in_json(result, amount)
+    
     return result
 
 iface = gr.Interface(
@@ -84,7 +87,7 @@ iface = gr.Interface(
     inputs=gr.Textbox(lines=2, placeholder="Enter a command..."),
     outputs=gr.Textbox(label="Extracted JSON Output"),
     title="T5 Fine-Tuned Command Extractor",
-    description="Extracts details in JSON format with exact amount preservation.",
+    description="Extracts details in JSON format and replaces the amount with the exact value from the input.",
 )
 
 if __name__ == "__main__":