Update app.py

app.py

@@ -1,5 +1,3 @@
-import json
-import re
 import torch
 import gradio as gr
 from transformers import T5Tokenizer, T5ForConditionalGeneration
@@ -9,43 +7,25 @@ model_name = "./t5-finetuned-final"
 tokenizer = T5Tokenizer.from_pretrained(model_name)
 model = T5ForConditionalGeneration.from_pretrained(model_name)
 
-# Move model to CPU explicitly
+# Move model to CPU
 device = torch.device("cpu")
 model.to(device)
 
-def extract_amount(input_text):
-    """Extracts the first numeric value (with decimals) from the input text."""
-    match = re.search(r'\b(\d+(?:\.\d+)?)\b', input_text)
-    return match.group(1) if match else None  # Keep as string
-
-def merge_json_with_amount(model_output, amount):
-    """Ensures the extracted amount is inserted into the model's output JSON."""
-    try:
-        data = json.loads(model_output)  # Try parsing the model output
-        if amount is not None:
-            data["amount"] = amount  # Replace model's amount with extracted amount (as string)
-        return json.dumps(data, ensure_ascii=False)
-    except json.JSONDecodeError:
-        return model_output  # Return raw model output if it's not valid JSON
-
 def generate_command(input_command):
-    """Processes the input, extracts the amount manually, and runs the model."""
-    amount = extract_amount(input_command)  # Extract amount
-
+    """Runs the model and returns its output as plain text."""
     input_ids = tokenizer("extract: " + input_command, return_tensors="pt").input_ids.to(device)
-    
-    output_ids = model.generate(input_ids, max_length=128, num_beams=1, early_stopping=True)  # Faster decoding
+
+    output_ids = model.generate(input_ids, max_length=128, num_beams=1, early_stopping=True)  # Simple decoding
     model_output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
-    result = merge_json_with_amount(model_output, amount)
-    return result
+    return model_output  # Return the model's raw output
 
 iface = gr.Interface(
     fn=generate_command,
     inputs=gr.Textbox(lines=2, placeholder="Enter a command..."),
-    outputs=gr.Textbox(label="Extracted JSON Output"),
-    title="T5 Command Extractor",
-    description="The model extracts action, currency, and recipient. The amount is manually extracted."
+    outputs=gr.Textbox(label="Model Output"),
+    title="T5 Model Output Test",
+    description="This version only shows the model's raw output."
 )
 
 if __name__ == "__main__":