import spaces
import gradio as gr
from wtpsplit import SaT
import json

# 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(duration=59)
def segment_text(input_text):
    # Use the SaT model to split the input text into sentences
    sentences = sat.split(input_text)
    
    # Create a JSON object where each sentence is a separate item
    json_output = json.dumps({"segments": sentences}, indent=2)
    
    return json_output

# Create the Gradio interface
iface = gr.Interface(
    fn=segment_text,
    inputs=gr.Textbox(lines=5, label="Input Text"),
    outputs=gr.JSON(label="Segmented Text (JSON)"),
    title="Text Segmentation with SaT",
    description="This app uses the SaT (Segment any Text) model to split input text into sentences and return the result as JSON. All credits to the respective author(s). Github: https://github.com/segment-any-text/wtpsplit/tree/main",
    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()