import gradio as gr
from transformers import pipeline
from span_marker import SpanMarkerModel

# Define the function
def function(messages):
    # Load the pre-trained model from Hugging Face
    model = SpanMarkerModel.from_pretrained("tomaarsen/span-marker-bert-base-acronyms")
    
    # Ensure the input is valid
    if messages:
        # Run inference on the input messages
        out = model.predict(messages)
        return out
    else:
        return "Please provide valid text input."

# Set up the Gradio interface
demo = gr.ChatInterface(
    fn=function,
    type="messages",
    examples=["Hello, I'm curious about abbreviations like NASA.", "Hola, ¿qué significa AI?", "Merhaba, kısaltmalar hakkında bilgi verir misiniz?"],
    title="Abbreviation Detector"
)

# Launch the app
if __name__ == "__main__":
    demo.launch()