Initial SQL-only 1C adapter baseline

scripts/train_1c_lora.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import yaml
+
+from common import ROOT, localize_workspace_path
+
+DEFAULT_CONFIG = ROOT / "plugins" / "1c" / "training" / "configs" / "qwen3-coder-30b-a3b-lora.yaml"
+
+
+def fail_missing_dependencies(exc: Exception) -> None:
+    raise SystemExit(
+        "Missing training dependencies. Install `requirements-training.txt` "
+        "or run inside the GPU training container. Original error: "
+        f"{type(exc).__name__}: {exc}"
+    )
+
+
+def load_config(path: Path) -> dict:
+    with path.open("r", encoding="utf-8") as handle:
+        data = yaml.safe_load(handle)
+    if not isinstance(data, dict):
+        raise ValueError(f"{path} must contain a YAML mapping")
+    return data
+
+
+def load_messages(path: Path) -> list[dict]:
+    rows = []
+    with path.open("r", encoding="utf-8") as handle:
+        for line in handle:
+            line = line.strip()
+            if line:
+                rows.append(json.loads(line))
+    return rows
+
+
+def build_training_texts(rows: list[dict], tokenizer) -> list[str]:
+    texts = []
+    for row in rows:
+        messages = row["messages"]
+        if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template:
+            text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
+        else:
+            text = "\n".join(f"{msg['role']}: {msg['content']}" for msg in messages)
+        texts.append(text)
+    return texts
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Train the 1C LoRA adapter.")
+    parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG)
+    parser.add_argument("--dry-run", action="store_true")
+    args = parser.parse_args()
+
+    config = load_config(args.config)
+    base_model_path = str(localize_workspace_path(config["base_model_path"]))
+    dataset_path = localize_workspace_path(config["dataset_path"])
+    output_dir = str(localize_workspace_path(config["output_dir"]))
+
+    rows = load_messages(dataset_path)
+    if args.dry_run:
+        print(f"Config: {args.config}")
+        print(f"Base model: {base_model_path}")
+        print(f"Dataset: {dataset_path} ({len(rows)} record(s))")
+        print(f"Output: {output_dir}")
+        return 0
+
+    try:
+        import torch
+        from datasets import Dataset
+        from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+        from transformers import (
+            AutoModelForCausalLM,
+            AutoTokenizer,
+            BitsAndBytesConfig,
+            DataCollatorForLanguageModeling,
+            Trainer,
+            TrainingArguments,
+        )
+    except Exception as exc:  # pragma: no cover
+        fail_missing_dependencies(exc)
+
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required for this training job.")
+
+    tokenizer = AutoTokenizer.from_pretrained(base_model_path, trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    qconf = config.get("quantization") or {}
+    quantization_config = None
+    if qconf.get("load_in_4bit"):
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type=qconf.get("bnb_4bit_quant_type", "nf4"),
+            bnb_4bit_use_double_quant=bool(qconf.get("bnb_4bit_use_double_quant", True)),
+            bnb_4bit_compute_dtype=torch.bfloat16,
+        )
+
+    model = AutoModelForCausalLM.from_pretrained(
+        base_model_path,
+        trust_remote_code=True,
+        device_map="auto",
+        quantization_config=quantization_config,
+    )
+    if quantization_config is not None:
+        model = prepare_model_for_kbit_training(model)
+
+    lora = config["lora"]
+    peft_config = LoraConfig(
+        r=int(lora["r"]),
+        lora_alpha=int(lora["lora_alpha"]),
+        lora_dropout=float(lora["lora_dropout"]),
+        bias=lora.get("bias", "none"),
+        task_type="CAUSAL_LM",
+        target_modules=lora["target_modules"],
+    )
+    model = get_peft_model(model, peft_config)
+
+    texts = build_training_texts(rows, tokenizer)
+    dataset = Dataset.from_dict({"text": texts})
+
+    def tokenize(batch):
+        return tokenizer(
+            batch["text"],
+            truncation=True,
+            max_length=int(config["max_seq_length"]),
+            padding=False,
+        )
+
+    tokenized = dataset.map(tokenize, batched=True, remove_columns=["text"])
+    collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+    train = config["train"]
+    train_args = TrainingArguments(
+        output_dir=output_dir,
+        num_train_epochs=float(train["num_train_epochs"]),
+        per_device_train_batch_size=int(train["per_device_train_batch_size"]),
+        gradient_accumulation_steps=int(train["gradient_accumulation_steps"]),
+        learning_rate=float(train["learning_rate"]),
+        warmup_ratio=float(train["warmup_ratio"]),
+        logging_steps=int(train["logging_steps"]),
+        save_strategy=train.get("save_strategy", "epoch"),
+        save_total_limit=int(train.get("save_total_limit", 1)),
+        bf16=bool(train.get("bf16", False)),
+        fp16=bool(train.get("fp16", False)),
+        report_to=[],
+    )
+    trainer = Trainer(model=model, args=train_args, train_dataset=tokenized, data_collator=collator)
+    trainer.train()
+    trainer.save_model(output_dir)
+    tokenizer.save_pretrained(output_dir)
+    print(f"Saved 1C LoRA adapter to {output_dir}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())