import spaces
import gradio as gr
from wtpsplit import SaT

# Initialize the SaT model
# We use 'sat-3l-sm' as it's recommended for general sentence segmentation tasks
# and offers a good balance between speed and performance
sat = SaT("sat-3l-sm")
sat.half().to("cuda")

@spaces.GPU
def segment_text(input_text):
    # Use the SaT model to split the input text into sentences
    sentences = sat.split(input_text)
    
    # Join the sentences with newlines for better readability in the output
    return "\n".join(sentences)

# Create the Gradio interface
iface = gr.Interface(
    fn=segment_text,
    inputs=gr.Textbox(lines=5, label="Input Text"),
    outputs=gr.Textbox(lines=10, label="Segmented Text"),
    title="Text Segmentation with SaT",
    description="This app uses the SaT (Segment any Text) model to split input text into sentences.",
    examples=[
        ["This is a test This is another test."],
        ["Hello this is a test But this is different now Now the next one starts looool"],
        ["The quick brown fox jumps over the lazy dog It was the best of times, it was the worst of times"],
    ]
)

# Launch the app
iface.launch()