app.py

import gradio as gr
import spaces
import os
import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer

# Nanbeige base model
MODEL_ID = "Nanbeige/Nanbeige4.1-3B"
DATASET_ID = "LvcidPsyche/openclaw-training-data-compressed"

@spaces.GPU(duration=1300) # Request ~21 minutes of ZeroGPU time
def train_model(hub_token, output_repo):
    if not hub_token or not output_repo:
        return "Please provide both HF Token and Output Repo Name."
    
    try:
        gr.Info("Loading Dataset...")
        # Add token to load_dataset since the dataset is private
        dataset = load_dataset(DATASET_ID, split="train", token=hub_token)

        gr.Info("Loading Tokenizer...")
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True, token=hub_token)
        tokenizer.pad_token = tokenizer.eos_token
        
        from transformers import BitsAndBytesConfig
        
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16
        )
        
        gr.Info("Loading Model in 4-bit...")
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID, 
            trust_remote_code=True, 
            quantization_config=bnb_config,
            device_map="auto",
            token=hub_token
        )
        model = prepare_model_for_kbit_training(model)
        
        peft_config = LoraConfig(
            r=16,
            lora_alpha=32,
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
            lora_dropout=0.05,
            bias="none",
            task_type="CAUSAL_LM"
        )
        # SFTTrainer will automatically apply this peft_config to the model.
        # Do not wrap the model manually with get_peft_model().

        def formatting_func(example):
            if "text" in example:
                return example["text"]
            elif "prompt" in example and "completion" in example:
                return f"Human: {example['prompt']}\n\nAssistant: {example['completion']}"
            elif "messages" in example:
                text = ""
                for msg in example["messages"]:
                    text += f"{msg['role'].capitalize()}: {msg['content']}\n"
                return text
            else:
                return str(example)

        from trl import SFTConfig

        training_args = SFTConfig(
            output_dir="./results",
            per_device_train_batch_size=2,
            gradient_accumulation_steps=4,
            learning_rate=2e-4,
            logging_steps=10,
            max_steps=200, # Demo steps to fit within 25 minutes
            push_to_hub=True,
            hub_model_id=output_repo,
            hub_token=hub_token,
            fp16=False,
            bf16=True,
            optim="paged_adamw_8bit",
            max_length=1024,
            dataset_text_field="text" if "text" in dataset.column_names else None
        )

        gr.Info("Starting Training on ZeroGPU...")
        trainer = SFTTrainer(
            model=model,
            train_dataset=dataset,
            peft_config=peft_config,
            processing_class=tokenizer,
            args=training_args,
            formatting_func=formatting_func if "text" not in dataset.column_names else None,
        )
        
        trainer.train()
        
        gr.Info("Pushing to Hub...")
        trainer.push_to_hub()
        
        return f"Training Complete! Model pushed to https://huggingface.co/{output_repo}"
    
    except Exception as e:
        return f"Error during training: {str(e)}"

with gr.Blocks(title="Nanbeige 3B Fine-Tuner", theme=gr.themes.Base()) as demo:
    gr.Markdown(f"# 🚀 Nanbeige 4.1 3B Fine-Tuning on ZeroGPU")
    gr.Markdown(f"This space fine-tunes `{MODEL_ID}` on your dataset `{DATASET_ID}` using LoRA and 4-bit quantization, then pushes it to your Hugging Face account. **ZeroGPU enabled for up to 25 minutes.**")
    
    with gr.Row():
        hf_token = gr.Textbox(label="HuggingFace Write Token (required to push model)", type="password", placeholder="hf_...")
        out_repo = gr.Textbox(label="Output Repo Name", placeholder="LvcidPsyche/nanbeige-finetuned")
        
    train_btn = gr.Button("Start Fine-Tuning", variant="primary")
    output_log = gr.Textbox(label="Status / Output", interactive=False)
    
    train_btn.click(fn=train_model, inputs=[hf_token, out_repo], outputs=output_log)

if __name__ == "__main__":
    demo.launch()