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

# Force CPU-only mode
os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ["BITSANDBYTES_NOWELCOME"] = "1"

# Configure logging
logging.basicConfig(level=logging.INFO)

def train():
    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,
            load_in_4bit=False  # Disable quantization
        )

        # Add padding token
        tokenizer.pad_token = tokenizer.eos_token

        # Load sample dataset
        dataset = load_dataset("wikitext", "wikitext-2-raw-v1")

        # Tokenization function
        def tokenize_function(examples):
            return tokenizer(
                examples["text"],
                padding="max_length",
                truncation=True,
                max_length=256,
                return_tensors="pt",
            )

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

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

        # Training arguments
        training_args = TrainingArguments(
            output_dir="./results",
            per_device_train_batch_size=2,
            per_device_eval_batch_size=2,
            num_train_epochs=1,  # Reduced for testing
            logging_dir="./logs",
            fp16=False,
            bf16=False,
            use_cpu=True  # Explicit CPU usage
        )

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

        # Start training
        logging.info("Starting training...")
        trainer.train()
        logging.info("Training completed!")
        
        return "✅ Training successful! Model saved."

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

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Phi-2 CPU Training")
    start_btn = gr.Button("Start Training")
    output = gr.Textbox()
    
    start_btn.click(
        fn=train,
        outputs=output
    )

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