import gradio as gr
import json
from transformers import pipeline

# Initialize the text generation pipeline
generator = pipeline('text-generation', model='gpt2')


def adjust_tone(text, concise, casual):
    tones = [
        {"tone": "concise", "weight": concise},
        {"tone": "casual", "weight": casual},
        {"tone": "professional", "weight": 1 - casual},
        {"tone": "expanded", "weight": 1 - concise}
    ]
    tones = sorted(tones, key=lambda x: x['weight'], reverse=True)[:2]

    tone_prompt = " and ".join([f"{t['tone']} (weight: {t['weight']:.2f})" for t in tones])

    prompt = f"Rewrite the following text to match these tones: {tone_prompt}. Text: {text}"

    result = generator(prompt, max_length=100, num_return_sequences=1)
    return result[0]['generated_text']


# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Tone Adjuster")

    input_text = gr.Textbox(label="Input Text")

    with gr.Row():
        concise_slider = gr.Slider(minimum=0, maximum=1, value=0.5, label="Concise vs Expanded")
        casual_slider = gr.Slider(minimum=0, maximum=1, value=0.5, label="Casual vs Professional")

    output_text = gr.Textbox(label="Adjusted Text")

    adjust_btn = gr.Button("Adjust Tone")

    adjust_btn.click(
        adjust_tone,
        inputs=[input_text, concise_slider, casual_slider],
        outputs=output_text
    )

demo.launch()