import gradio as gr

def train_model(epochs, batch_size, learning_rate):
    return f"Training completed!\nEpochs: {epochs}\nBatch Size: {batch_size}\nLearning Rate: {learning_rate}"

with gr.Blocks(title="DistilBERT Arabic Sentiment Training") as demo:
    gr.Markdown("# DistilBERT Arabic Sentiment Training")
    gr.Markdown("Fine-tune DistilBERT on Arabic sentiment analysis (Saudi dialect)")
    
    gr.Markdown("### Model Information:")
    gr.Markdown("- **Base Model**: distilbert-base-multilingual-cased (67M parameters)")
    gr.Markdown("- **Task**: Text Classification (Multilingual)")
    gr.Markdown("- **Dataset**: arbml/Arabic_Sentiment_Twitter_Corpus (58.8k examples)")
    gr.Markdown("- **Language**: Arabic (Saudi & Gulf dialects)")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Training Settings")
            epochs = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="Epochs")
            batch_size = gr.Slider(minimum=8, maximum=64, value=32, step=8, label="Batch Size")
            learning_rate = gr.Slider(minimum=1e-5, maximum=1e-3, value=2e-5, step=1e-5, label="Learning Rate")
        
        with gr.Column():
            gr.Markdown("### Training Status")
            output_text = gr.Textbox(label="Output", lines=10, interactive=False)
    
    train_button = gr.Button("Start Training", variant="primary", size="lg")
    train_button.click(
        fn=train_model,
        inputs=[epochs, batch_size, learning_rate],
        outputs=output_text
    )
    
    gr.Markdown("### Training Details:")
    gr.Markdown("- **Hardware**: Free GPU (Hugging Face Spaces)")
    gr.Markdown("- **Expected Time**: 5-10 minutes (GPU) or 15-20 minutes (CPU)")
    gr.Markdown("- **Output Directory**: ./results")
    gr.Markdown("- **Usage**: Arabic text only")

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