import gradio as gr

from chonky import ParagraphSplitter


def run(text, model_id):
    splitter = ParagraphSplitter(model_id=model_id)
    return "\n\n---\n\n".join(splitter(text))

with gr.Blocks() as demo:
    gr.Markdown("# Semantic Chunking Demo\n  **Note**: This Space runs on CPU only, so input is limited to max. 50000 characters.")
    model = gr.Dropdown(label="Choose model", value="mamei16/chonky_distilbert_base_uncased_1.1",
                        choices=["mamei16/chonky_distilbert_base_uncased_1.1", "mamei16/chonky_mdistilbert-base-english-cased",
                                 "mirth/chonky_distilbert_base_uncased_1", "mirth/chonky_modernbert_base_1"])
    button = gr.Button("Run", variant="primary")
    text = gr.Textbox(label='Input Text', max_length=50000)
    gr.Markdown("## Result chunks:")
    chunks = gr.Markdown("---")

    button.click(run, [text, model], chunks)


if __name__ == "__main__":
    demo.queue(max_size=20)
    demo.launch()