import gradio as gr
import spaces
import os, subprocess

@spaces.GPU(duration=600)
def train():
    log = []
    subprocess.run(['pip','install','-q','unsloth','datasets','trl','peft','accelerate','bitsandbytes'],capture_output=True)
    log.append('Packages installed')
    subprocess.run(['wget','-q','https://huggingface.co/datasets/yace222/weval-brain-dataset/resolve/main/train_chatml.jsonl','-O','/tmp/train.jsonl'],capture_output=True)
    with open('/tmp/train.jsonl') as f: count=sum(1 for _ in f)
    log.append(f'Dataset: {count} pairs')
    try:
        from unsloth import FastLanguageModel
        import torch
        model,tok = FastLanguageModel.from_pretrained(model_name='unsloth/Qwen2.5-7B-Instruct-bnb-4bit',max_seq_length=2048,load_in_4bit=True)
        model = FastLanguageModel.get_peft_model(model,r=16,target_modules=['q_proj','k_proj','v_proj','o_proj','gate_proj','up_proj','down_proj'],lora_alpha=16,lora_dropout=0,bias='none',use_gradient_checkpointing='unsloth')
        log.append('Model+LoRA ready')
        from datasets import load_dataset
        ds = load_dataset('json',data_files='/tmp/train.jsonl',split='train')
        def fmt(ex): return {'text':tok.apply_chat_template(ex['messages'],tokenize=False,add_generation_prompt=False)}
        ds = ds.map(fmt)
        from trl import SFTTrainer
        from transformers import TrainingArguments
        trainer = SFTTrainer(model=model,tokenizer=tok,train_dataset=ds,dataset_text_field='text',max_seq_length=2048,args=TrainingArguments(per_device_train_batch_size=2,gradient_accumulation_steps=4,warmup_steps=5,num_train_epochs=3,learning_rate=2e-4,fp16=not torch.cuda.is_bf16_supported(),bf16=torch.cuda.is_bf16_supported(),logging_steps=1,output_dir='outputs',optim='adamw_8bit',seed=42))
        stats = trainer.train()
        log.append(f'Loss: {stats.training_loss:.4f}')
        model.save_pretrained_gguf('gguf',tok,quantization_method='q4_k_m')
        HF_TOKEN = os.environ.get('HF_TOKEN','')
        if HF_TOKEN:
            model.push_to_hub_gguf('yace222/weval-brain-v3-gguf',tok,quantization_method='q4_k_m',token=HF_TOKEN)
            log.append('Pushed to HF!')
    except Exception as e:
        log.append(f'ERROR: {e}')
    return chr(10).join(log)

demo = gr.Interface(fn=train,inputs=[],outputs=gr.Textbox(lines=20),title='WEVAL Brain AutoTrain')
demo.launch()