import gradio as gr
from humanize import paraphrase_text

with gr.Blocks() as demo:
    gr.Markdown("# Polygraf Writer")
    with gr.Row():
        with gr.Column(scale=0.7):
            gr.Markdown("## Enter a topic to write an article about:")
            input_topic = gr.Textbox(label="Topic")
            gr.Markdown("### Humanized article:")
            output_label = gr.HTML(label="Output")
            model_dropdown = gr.Radio(
                choices=[
                    "Base Model",
                    "Large Model",
                    "XL Model",
                    "XL Law Model",
                    "XL Marketing Model",
                    "XL Child Style Model",
                ],
                value="Large Model",
                label="Select Model Version",
            )
            process_button = gr.Button("Humanize Text")

        with gr.Column(scale=0.3):
                temperature_slider = gr.Slider(minimum=0.5, maximum=2.0, step=0.1, value=1.2, label="Temperature")
                gr.Markdown("Controls the randomness of the paraphrase. Higher values generate more varied text.")
                top_k_slider = gr.Slider(
                    minimum=0,
                    maximum=300,
                    step=25,
                    value=50,
                    label="Top k",
                )
                gr.Markdown("Limits the number of top tokens considered during generation.")
                repetition_penalty_slider = gr.Slider(
                    minimum=1.0,
                    maximum=2.0,
                    step=0.1,
                    value=1,
                    label="Repetition Penalty",
                )
                gr.Markdown("Penalizes repeated words to encourage diverse language use")
                length_penalty_slider = gr.Slider(
                    minimum=0.0,
                    maximum=2.0,
                    step=0.1,
                    value=1.0,
                    label="Length Penalty",
                )
                gr.Markdown("Penalizes shorter outputs.")
    
        process_button.click(
            fn=paraphrase_text,
            inputs=[
                input_text,
                model_dropdown,
                temperature_slider,
                repetition_penalty_slider,
                top_k_slider,
                length_penalty_slider,
            ],
            outputs=output_label,
        )

if __name__ == "__main__":
    demo.launch(demo.launch(server_name="0.0.0.0"))