import gradio as gr

# Load the model
interface = gr.load("models/alexue4/text-normalization-ru-new")

# Function to modify the behavior (example: limiting max tokens and setting examples)
def custom_function(input_text):
    # Custom processing here (e.g., limit max tokens)
    max_tokens = 250  # Set your desired max token limit
    if len(input_text.split()) > max_tokens:
        input_text = ' '.join(input_text.split()[:max_tokens])
    
    # Call the original function from the interface
    original_output = interface.predict(input_text)
    return original_output

examples = [
    ["в 2006-2010 гг. Вася учился в МГУ"],
    ["я купил iphone 10X за 14990 руб без 3-x часов полдень и т.д."]
]

gr.Interface(
    fn=custom_function,
    inputs="text",
    outputs="text",
    examples=examples
).launch()