import gradio as gr
from transformers import pipeline

# Load a Hugging Face model for rewriting
rewrite_pipeline = pipeline("text2text-generation", model="t5-small")
classification_pipeline = pipeline("text-classification", model="textattack/bert-base-uncased-imdb")

def rewrite_and_analyze(input_text, temperature):
    # Generate rewritten text
    rewritten = rewrite_pipeline(input_text, temperature=temperature, max_length=200)
    rewritten_text = rewritten[0]['generated_text']
    
    # Analyze classification probabilities
    analysis = classification_pipeline(rewritten_text, return_all_scores=True)[0]
    class_probs = {item['label']: round(item['score'], 3) for item in analysis}
    
    return rewritten_text, class_probs

# Gradio interface
inputs = [
    gr.Textbox(lines=5, placeholder="Enter text here...", label="Input Text"),
    gr.Slider(0.1, 2.0, value=1.0, label="Temperature")
]
outputs = [
    gr.Textbox(label="Rewritten Text"),
    gr.JSON(label="Class Probabilities")
]

app = gr.Interface(
    fn=rewrite_and_analyze,
    inputs=inputs,
    outputs=outputs,
    title="Text Rewriter and Classifier"
)

app.launch()