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

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

def train():
    model = AutoModelForCausalLM.from_pretrained(
        "microsoft/phi-2",
        device_map="auto",
        trust_remote_code=True,
        load_in_4bit=False  # Disable quantization
    )
    
    training_args = TrainingArguments(
        output_dir="./results",
        per_device_train_batch_size=2,
        num_train_epochs=3,
        use_cpu=True,  # Explicit CPU usage
        fp16=False,
        bf16=False,
    )
    
    # Rest of training code...