train_abap.py

"""
Fine-tune Qwen2.5-Coder-7B-Instruct for ABAP development using SFT + LoRA.

Combines multiple ABAP instruction datasets into a unified conversational format,
then trains with TRL's SFTTrainer.

Based on:
- Low-resource PL fine-tuning insights (arxiv:2501.19085)
- Qwen2.5-Coder training patterns (arxiv:2409.12186)
- Reference: huggingface/trl examples/scripts/sft.py

Datasets:
- smjain/abap: 248 ABAP coding prompt/response pairs
- Kaballas/abap: 1,070 ABAP concept Q&A
- Arturs213/abap-code-sec-finetune: ~4K+ ABAP security analysis (all splits)

Run:
  pip install torch transformers trl peft datasets accelerate bitsandbytes
  python train_abap.py
"""

import os
import torch
from datasets import load_dataset, concatenate_datasets
from trl import SFTTrainer, SFTConfig
from peft import LoraConfig, TaskType

# ─── Configuration ───────────────────────────────────────────────────────────
MODEL_ID = "Qwen/Qwen2.5-Coder-7B-Instruct"
HUB_MODEL_ID = "SpaceArm/Qwen2.5-Coder-7B-ABAP"
OUTPUT_DIR = "./qwen25-coder-abap-lora"

SYSTEM_PROMPT = (
    "You are an expert SAP ABAP developer. You write clean, modern ABAP code "
    "following SAP best practices. You can help with ABAP reports, classes, "
    "function modules, ALV grids, internal tables, SELECT statements, BAPIs, "
    "RAP (RESTful ABAP Programming), CDS views, and all aspects of SAP development."
)


# ─── Dataset preparation ─────────────────────────────────────────────────────
def load_smjain_abap():
    """smjain/abap: ~248 prompt/response ABAP coding examples."""
    ds = load_dataset("smjain/abap", split="train")
    return ds.map(lambda x: {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": x["prompt"]},
        {"role": "assistant", "content": x["response"]},
    ]}, remove_columns=ds.column_names)


def load_kaballas_abap():
    """Kaballas/abap: ~1,070 Q&A about ABAP concepts and OOP patterns."""
    ds = load_dataset("Kaballas/abap", split="train")
    return ds.map(lambda x: {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": x["question"]},
        {"role": "assistant", "content": x["answer"]},
    ]}, remove_columns=ds.column_names)


def load_arturs_abap_sec():
    """Arturs213/abap-code-sec-finetune: ABAP security analysis (all splits)."""
    datasets = []
    for split_name in ["base", "expanded", "clear"]:
        ds = load_dataset("Arturs213/abap-code-sec-finetune", split=split_name)
        ds = ds.map(lambda x: {"messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": (x["Instruction"] + "\n\n" + x["Input"]) if x["Input"] else x["Instruction"]},
            {"role": "assistant", "content": x["Response"]},
        ]}, remove_columns=ds.column_names)
        datasets.append(ds)
    return concatenate_datasets(datasets)


def prepare_dataset():
    """Load and combine all ABAP instruction datasets."""
    print("=" * 60)
    print("Loading ABAP instruction datasets")
    print("=" * 60)

    print("\n[1/3] smjain/abap...")
    ds1 = load_smjain_abap()
    print(f"  -> {len(ds1)} examples")

    print("[2/3] Kaballas/abap...")
    ds2 = load_kaballas_abap()
    print(f"  -> {len(ds2)} examples")

    print("[3/3] Arturs213/abap-code-sec-finetune (all splits)...")
    ds3 = load_arturs_abap_sec()
    print(f"  -> {len(ds3)} examples")

    # Combine and shuffle
    combined = concatenate_datasets([ds1, ds2, ds3])
    combined = combined.shuffle(seed=42)

    print(f"\n{'=' * 60}")
    print(f"Total training examples: {len(combined)}")
    print(f"{'=' * 60}")

    return combined


# ─── Main training ───────────────────────────────────────────────────────────
def main():
    # Prepare data
    full_dataset = prepare_dataset()

    # 95/5 train/eval split
    split = full_dataset.train_test_split(test_size=0.05, seed=42)
    train_dataset = split["train"]
    eval_dataset = split["test"]
    print(f"\nTrain: {len(train_dataset)} | Eval: {len(eval_dataset)}")

    # LoRA config - rank 32, targeting all attention + MLP projections
    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=32,
        lora_alpha=64,
        lora_dropout=0.05,
        target_modules=[
            "q_proj", "k_proj", "v_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj",
        ],
        bias="none",
    )

    # SFT training config
    sft_config = SFTConfig(
        output_dir=OUTPUT_DIR,
        hub_model_id=HUB_MODEL_ID,

        # Schedule
        num_train_epochs=3,
        per_device_train_batch_size=2,
        per_device_eval_batch_size=2,
        gradient_accumulation_steps=8,  # effective batch = 16
        learning_rate=2e-4,
        lr_scheduler_type="cosine",
        warmup_steps=30,
        weight_decay=0.01,

        # SFT
        max_length=2048,
        packing=False,
        dataset_num_proc=4,

        # Memory / precision
        bf16=True,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},

        # Logging - plain text, no tqdm
        logging_steps=5,
        logging_first_step=True,
        disable_tqdm=True,
        logging_strategy="steps",

        # Eval & save
        eval_strategy="steps",
        eval_steps=100,
        save_strategy="steps",
        save_steps=100,
        save_total_limit=3,
        load_best_model_at_end=True,
        metric_for_best_model="eval_loss",

        # Hub
        push_to_hub=True,
        hub_strategy="every_save",

        # Tracking
        report_to="none",
    )

    # Trainer
    trainer = SFTTrainer(
        model=MODEL_ID,
        args=sft_config,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        peft_config=lora_config,
    )

    # Train
    print("\n Starting ABAP SFT training...")
    print(f"   Model: {MODEL_ID}")
    print(f"   LoRA rank: {lora_config.r}, alpha: {lora_config.lora_alpha}")
    print(f"   Effective batch size: {sft_config.per_device_train_batch_size * sft_config.gradient_accumulation_steps}")
    print(f"   Learning rate: {sft_config.learning_rate}")
    print(f"   Epochs: {sft_config.num_train_epochs}")
    print(f"   Max length: {sft_config.max_length}")

    train_result = trainer.train()

    # Final metrics
    metrics = train_result.metrics
    print(f"\n Training complete!")
    for k, v in metrics.items():
        print(f"   {k}: {v}")

    # Save & push
    print("\n Saving model and pushing to Hub...")
    trainer.save_model()
    trainer.push_to_hub(commit_message="ABAP fine-tuned Qwen2.5-Coder-7B with LoRA")

    print(f"\n Done! Model: https://huggingface.co/{HUB_MODEL_ID}")


if __name__ == "__main__":
    main()