import gradio as gr
from span_marker import SpanMarkerModel

# Define the function
def function(text):
    # Load the pre-trained model from Hugging Face
    model = SpanMarkerModel.from_pretrained("tomaarsen/span-marker-bert-base-acronyms")
    
    # Ensure the input is valid
    if text:
        # Run inference on the input text
        out = model.predict([text])  # Wrap the text in a list as the model expects a list of strings
        return out
    else:
        return "Please provide valid text input."

# Set up the Gradio interface
demo = gr.Interface(
    fn=function,
    inputs=gr.Textbox(lines=5, placeholder="Enter your text here...", label="Input Text"),
    outputs="json",
    title="Abbreviation Detector",
    examples=[
        "NASA is an abbreviation for National Aeronautics and Space Administration.",
        "AI stands for Artificial Intelligence.",
        "What does HTTP mean?"
    ]
)

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