App.py

# =============================================================================
# Full-Parameter SFT: Qwen/Qwen3.5-4B-Base on Bc-AI/SFT-Ultra
# On-the-fly streaming | Bad row filtering | HF Hub upload
# =============================================================================
# Requirements:

#   pip install torch transformers datasets trl accelerate huggingface_hub

import os
import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTTrainer, SFTConfig
from huggingface_hub import HfApi, login

# =============================================================================
# 0. CONFIG — Edit these as needed
# =============================================================================
MODEL_ID        = "Qwen/Qwen3.5-4B"
DATASET_ID      = "Bc-AI/SFT-Ultra"
OUTPUT_DIR      = "./qwen3.5-4b-full-sft"

# HF Hub — paste your token when prompted at runtime
HF_REPO_ID      = "Bc-AI/qwen3.5-4b-sft"   # ← change this
HF_TOKEN        = ""    # Leave None — you will be prompted to paste it below

# Sequence
MAX_SEQ_LENGTH  = 2048

# Training hyperparameters
NUM_TRAIN_EPOCHS              = 1
PER_DEVICE_TRAIN_BATCH_SIZE   = 2
GRADIENT_ACCUMULATION_STEPS   = 8
LEARNING_RATE                 = 1e-5
WEIGHT_DECAY                  = 0.01
WARMUP_RATIO                  = 0.05
LR_SCHEDULER                  = "cosine"
MAX_GRAD_NORM                 = 1.0

# Logging & saving
LOGGING_STEPS    = 10
SAVE_STEPS       = 500
SAVE_TOTAL_LIMIT = 3

# Precision — bf16 on Ampere+ (A100, 3090, 4090)
#             set bf16=False fp16=True on older GPUs (V100, T4)
USE_BF16 = True
USE_FP16 = False

# Streaming = on-the-fly download, no full disk pre-cache
STREAM_DATASET = True

SEED = 42

# =============================================================================
# 1. HF HUB LOGIN  — paste token here at runtime
# =============================================================================
print("=" * 60)
print("  Hugging Face Hub Login")
print("=" * 60)

if HF_TOKEN is None:
    HF_TOKEN = input("  Paste your HF token (hf_…): ").strip()

login(token=HF_TOKEN, add_to_git_credential=False)
print("  ✅ Logged in successfully.\n")

# =============================================================================
# 2. TOKENIZER
# =============================================================================
print(f"[1/4] Loading tokenizer: {MODEL_ID}")

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

tokenizer.padding_side = "right"

# =============================================================================
# 3. MODEL — full bf16, no quantisation, no adapter
# =============================================================================
print(f"[2/4] Loading full model: {MODEL_ID}")

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16 if USE_BF16 else torch.float32,
    device_map="auto",
)

model.config.use_cache = False

for param in model.parameters():
    param.requires_grad = True

total_params = sum(p.numel() for p in model.parameters())
print(f"   Trainable parameters: {total_params:,}  ({total_params / 1e9:.2f}B)")

# =============================================================================
# 4. DATASET — streaming + bad row filtering + on-the-fly tokenisation
# =============================================================================
print(f"[3/4] Loading dataset: {DATASET_ID}  (streaming={STREAM_DATASET})")

raw_dataset = load_dataset(
    DATASET_ID,
    split="train",
    streaming=STREAM_DATASET,
    trust_remote_code=True,
)

# ── Bad row validator ────────────────────────────────────────────────────────
# Catches every known failure mode so no single row can crash training

REQUIRED_ROLES = {"user", "assistant"}   # Minimum roles a valid convo must have

def is_valid_row(example):
    """
    Returns True only if the row is safe to train on.
    Filters out:
      - Missing / non-list messages field
      - Empty message list
      - Messages with missing role or content keys
      - Messages where role or content is not a string
      - Messages where content is an empty / whitespace-only string
      - Conversations missing at least one user AND one assistant turn
      - Rows where the entire rendered text would be empty
    """
    try:
        messages = example.get("messages", None)

        # Must exist and be a non-empty list
        if not isinstance(messages, list) or len(messages) == 0:
            return False

        seen_roles = set()
        for msg in messages:
            # Each message must be a dict
            if not isinstance(msg, dict):
                return False

            role    = msg.get("role",    None)
            content = msg.get("content", None)

            # role and content must be non-empty strings
            if not isinstance(role, str) or not role.strip():
                return False
            if not isinstance(content, str) or not content.strip():
                return False

            seen_roles.add(role.strip().lower())

        # Must have at least one user turn and one assistant turn
        if not REQUIRED_ROLES.issubset(seen_roles):
            return False

        return True

    except Exception:
        # Catch-all: any unexpected structure is silently dropped
        return False


# ── Safe formatter ───────────────────────────────────────────────────────────

def format_messages(example):
    """
    Applies the Qwen3.5 chat template.
    Wrapped in try/except so any template rendering failure is handled
    gracefully — the row is marked with an empty text field and later dropped.
    """
    try:
        text = tokenizer.apply_chat_template(
            example["messages"],
            tokenize=False,
            add_generation_prompt=False,
        )
        # Final safety: rendered text must be non-trivial
        if not text or not text.strip():
            return {"text": ""}
        return {"text": text}
    except Exception:
        return {"text": ""}


def is_non_empty_text(example):
    """Drop any row where formatting produced an empty string."""
    text = example.get("text", "")
    return isinstance(text, str) and len(text.strip()) > 0


# ── Apply pipeline ───────────────────────────────────────────────────────────
print("   Step 1 — Filtering malformed rows …")
clean_dataset = raw_dataset.filter(is_valid_row)

print("   Step 2 — Applying chat template on the fly …")
formatted_dataset = clean_dataset.map(format_messages)

print("   Step 3 — Dropping any rows with empty rendered text …")
formatted_dataset = formatted_dataset.filter(is_non_empty_text)

print("   ✅ Dataset pipeline ready.\n")

# =============================================================================
# 5. TRAINER
# =============================================================================
print("[4/4] Configuring SFTTrainer …")

sft_config = SFTConfig(
    # ── Output ───────────────────────────────────────────────────────────────
    output_dir=OUTPUT_DIR,

    # ── Sequence ─────────────────────────────────────────────────────────────
    max_seq_length=MAX_SEQ_LENGTH,

    # ── Training schedule ────────────────────────────────────────────────────
    num_train_epochs=NUM_TRAIN_EPOCHS,
    per_device_train_batch_size=PER_DEVICE_TRAIN_BATCH_SIZE,
    gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,
    learning_rate=LEARNING_RATE,
    weight_decay=WEIGHT_DECAY,
    warmup_ratio=WARMUP_RATIO,
    lr_scheduler_type=LR_SCHEDULER,
    max_grad_norm=MAX_GRAD_NORM,

    # ── Optimizer ────────────────────────────────────────────────────────────
    optim="adamw_torch_fused",

    # ── Precision ────────────────────────────────────────────────────────────
    bf16=USE_BF16,
    fp16=USE_FP16,

    # ── Gradient checkpointing ───────────────────────────────────────────────
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},

    # ── Loss ─────────────────────────────────────────────────────────────────
    completion_only_loss=True,

    # ── Dataset ──────────────────────────────────────────────────────────────
    dataset_text_field="text",
    dataset_num_proc=1,             # Must be 1 for IterableDataset (streaming)

    # ── Logging & checkpointing ──────────────────────────────────────────────
    logging_steps=LOGGING_STEPS,
    save_steps=SAVE_STEPS,
    save_total_limit=SAVE_TOTAL_LIMIT,
    report_to="none",               # Swap to "wandb" or "tensorboard" if needed

    # ── Misc ─────────────────────────────────────────────────────────────────
    seed=SEED,
    remove_unused_columns=True,
)

trainer = SFTTrainer(
    model=model,
    args=sft_config,
    train_dataset=formatted_dataset,
    tokenizer=tokenizer,
)

# =============================================================================
# 6. TRAIN
# =============================================================================
print("\n🚀  Starting full fine-tuning …\n")
trainer.train()

# =============================================================================
# 7. SAVE LOCALLY
# =============================================================================
print(f"\n💾  Saving full model + tokenizer to: {OUTPUT_DIR}")
trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)
print("   ✅ Local save complete.\n")

# =============================================================================
# 8. PUSH TO HF HUB
# =============================================================================
print(f"☁️   Uploading to Hugging Face Hub: {HF_REPO_ID}")
print("    (This may take a while depending on your upload speed …)\n")

try:
    # Push model
    model.push_to_hub(
        HF_REPO_ID,
        token=HF_TOKEN,
        commit_message="Full SFT — Qwen3.5-4B-Base on Bc-AI/SFT-Ultra",
        private=True,              # Set False if you want a public repo
    )

    # Push tokenizer
    tokenizer.push_to_hub(
        HF_REPO_ID,
        token=HF_TOKEN,
        commit_message="Add tokenizer",
    )

    # Push a minimal model card so the repo is well-documented
    api = HfApi()
    model_card = f"""---
language:
- en
license: apache-2.0
base_model: {MODEL_ID}
datasets:
- {DATASET_ID}
tags:
- full-fine-tune
- sft
- qwen3.5
---

# Qwen3.5-4B — Full SFT

- **Base model:** `{MODEL_ID}`
- **Dataset:** `{DATASET_ID}`
- **Training type:** Full parameter supervised fine-tuning (no LoRA)
- **Max sequence length:** {MAX_SEQ_LENGTH}
- **Epochs:** {NUM_TRAIN_EPOCHS}
- **Learning rate:** {LEARNING_RATE}
- **Precision:** {"bf16" if USE_BF16 else "fp16"}
"""
    api.upload_file(
        path_or_fileobj=model_card.encode("utf-8"),
        path_in_repo="README.md",
        repo_id=HF_REPO_ID,
        token=HF_TOKEN,
        commit_message="Add model card",
    )

    print(f"\n✅  Model successfully uploaded to: https://huggingface.co/{HF_REPO_ID}")

except Exception as e:
    print(f"\n❌  Upload failed: {e}")
    print(f"    Your model is still saved locally at: {OUTPUT_DIR}")
    print("    You can retry the upload manually with:")
    print(f"      model.push_to_hub('{HF_REPO_ID}')")
    print(f"      tokenizer.push_to_hub('{HF_REPO_ID}')")

print("\n✅  All done!")