import torch
import gradio as gr
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling
)
from datasets import load_dataset
import logging
import sys

# Configure logging
logging.basicConfig(stream=sys.stdout, level=logging.INFO)

def train(dataset_name: str, dataset_config: str = None):
    try:
        # Load model and tokenizer
        model_name = "microsoft/phi-2"
        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cpu", trust_remote_code=True)

        # Add padding token
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token

        # Load dataset from Hugging Face Hub
        logging.info(f"Loading dataset: {eswardivi/medical_qa} (config: {dataset_config})")
        dataset = load_dataset(
            dataset_name,
            dataset_config,  # Optional config (e.g., language for Common Voice)
            split="train+validation",  # Combine splits
            trust_remote_code=True  # Required for some datasets
        )

        # Split into train/validation
        dataset = dataset.train_test_split(test_size=0.1, seed=42)

        # Tokenization function (adjust based on dataset columns)
        def tokenize_function(examples):
            return tokenizer(
                examples["text"],  # Replace "text" with your dataset's text column
                padding="max_length",
                truncation=True,
                max_length=256,
                return_tensors="pt",
            )

        tokenized_dataset = dataset.map(
            tokenize_function,
            batched=True,
            remove_columns=dataset["train"].column_names
        )

        # Data collator
        data_collator = DataCollatorForLanguageModeling(
            tokenizer=tokenizer,
            mlm=False
        )

        # Training arguments
        training_args = TrainingArguments(
            output_dir="./phi2-results",
            per_device_train_batch_size=2,
            per_device_eval_batch_size=2,
            num_train_epochs=3,
            logging_dir="./logs",
            logging_steps=10,
            fp16=False,
        )

        # Trainer
        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=tokenized_dataset["train"],
            eval_dataset=tokenized_dataset["test"],
            data_collator=data_collator,
        )

        # Start training
        logging.info("Training started...")
        trainer.train()
        trainer.save_model("./phi2-trained-model")
        logging.info("Training completed!")

        return "✅ Training succeeded! Model saved."

    except Exception as e:
        logging.error(f"Training failed: {str(e)}")
        return f"❌ Training failed: {str(e)}"

# Gradio UI with dataset input
with gr.Blocks(title="Phi-2 Training") as demo:
    gr.Markdown("# 🚀 Train Phi-2 with HF Hub Data")
    
    with gr.Row():
        dataset_name = gr.Textbox(label="Dataset Name", value="mozilla-foundation/common_voice_11_0")
        dataset_config = gr.Textbox(label="Dataset Config (optional)", value="en")
    
    start_btn = gr.Button("Start Training", variant="primary")
    status_output = gr.Textbox(label="Status", interactive=False)
    
    start_btn.click(
        fn=train,
        inputs=[dataset_name, dataset_config],
        outputs=status_output
    )

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