Update app.py

app.py

@@ -14,44 +14,34 @@ device = torch.device("cpu")
 model.to(device)
 
 def extract_amount(input_text):
-    """
-    Extracts the first number (with optional decimals) from the input text.
-    For example, in:
-      "Should I send 2659.53464 EUR to my wife today?"
-    it returns the string "2659.53464".
-    """
+    """Extracts the first numeric value (with decimals) from the input text."""
     match = re.search(r'\b(\d+(?:\.\d+)?)\b', input_text)
-    if match:
-        return match.group(1)
-    return None
+    return match.group(1) if match else None
 
 def merge_json_with_amount(model_output, amount):
-    """
-    Parses the model's JSON output and overrides the "amount" key
-    with the manually extracted value.
-    """
+    """Ensures the extracted amount is inserted into the model's output JSON."""
     try:
-        data = json.loads(model_output)
-    except Exception:
-        data = {}
+        data = json.loads(model_output)  # Try parsing the model output
+    except json.JSONDecodeError:
+        return model_output  # Return raw model output if it's not valid JSON
+
     if amount:
         try:
-            data["amount"] = float(amount)
-        except Exception:
-            data["amount"] = amount
+            data["amount"] = float(amount)  # Convert to float
+        except ValueError:
+            data["amount"] = amount  # Keep as string if conversion fails
+
     return json.dumps(data, ensure_ascii=False)
 
 def generate_command(input_command):
-    # Manually extract the amount from the input.
-    amount = extract_amount(input_command)
+    """Processes the input, extracts the amount manually, and runs the model."""
+    amount = extract_amount(input_command)  # Extract amount
+
+    input_ids = tokenizer("extract: " + input_command, return_tensors="pt").input_ids.to(device)
     
-    # Generate the JSON output from the model.
-    prompt = "extract: " + input_command
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
-    output_ids = model.generate(input_ids, max_length=128, num_beams=2, early_stopping=True)
+    output_ids = model.generate(input_ids, max_length=128, num_beams=1, early_stopping=True)  # Faster decoding
     model_output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    
-    # Merge the manually extracted amount into the model output.
+
     result = merge_json_with_amount(model_output, amount)
     return result
 
@@ -60,7 +50,7 @@ iface = gr.Interface(
     inputs=gr.Textbox(lines=2, placeholder="Enter a command..."),
     outputs=gr.Textbox(label="Extracted JSON Output"),
     title="T5 Command Extractor",
-    description="The model provides action, currency, and recipient. The amount is manually extracted from the input."
+    description="The model extracts action, currency, and recipient. The amount is manually extracted."
 )
 
 if __name__ == "__main__":