import re
import torch
import gradio as gr
from transformers import T5Tokenizer, T5ForConditionalGeneration

# 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 CPU
device = torch.device("cpu")
model.to(device)

def extract_amount(input_text):
    """Finds the first number in the text (supports both . and , as decimal separators)."""
    match = re.search(r'\b\d{1,15}([.,]\d{1,15})?\b', input_text)
    return match.group(0) if match else "Not found"  # Keep as string

def extract_variable(text, key):
    """Extracts a value following a key in the model output."""
    match = re.search(rf'"{key}"\s*:\s*"([^"]+)"', text)  # Matches "key": "value"
    return match.group(1) if match else "Not found"

def generate_command(input_command):
    """Runs the model, extracts amount, and processes output variables."""
    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)
    model_output = tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()

    # Extract amount manually
    extracted_amount = extract_amount(input_command)

    # Extract the rest of the values from the raw model output
    action = extract_variable(model_output, "action")
    currency = extract_variable(model_output, "currency")
    recipient = extract_variable(model_output, "recipient")

    return extracted_amount, action, currency, recipient

# Gradio Interface
iface = gr.Interface(
    fn=generate_command,
    inputs=gr.Textbox(lines=2, placeholder="Enter a command..."),
    outputs=[
        gr.Textbox(label="Amount"),  
        gr.Textbox(label="Action"),
        gr.Textbox(label="Currency"),
        gr.Textbox(label="Recipient"),
    ],
    title="Intent Recognition Project",
    description="Extracts amount separately and displays action, currency, and recipient from model output."
)

if __name__ == "__main__":
    iface.launch()