import gradio as gr
from datasets import load_dataset
from transformers import Trainer, TrainingArguments, AutoModelForSequenceClassification

def train_model():
    # Načti dataset
    dataset = load_dataset("imdb")  # Můžeš nahradit za jiný dataset podle potřeby
    
    # Načti model
    model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)
    
    # Nastavení trénovacích argumentů
    training_args = TrainingArguments(
        output_dir="./results",
        evaluation_strategy="epoch",
        learning_rate=2e-5,
        per_device_train_batch_size=16,
        per_device_eval_batch_size=64,
        num_train_epochs=3,
    )
    
    # Vytvoření trénovacího objektu
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset['train'],
        eval_dataset=dataset['test'],
    )
    
    # Trénování modelu
    trainer.train()

    return "Model has been trained!"

# Gradio rozhraní pro spuštění trénování
demo = gr.Interface(fn=train_model, inputs=[], outputs="text", live=True)

# Spuštění aplikace
demo.launch()