| import argparse |
| import json |
| import os |
| import random |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Dict, List |
|
|
| import torch |
| from datasets import load_dataset |
| from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training |
| from transformers import ( |
| AutoModelForCausalLM, |
| AutoTokenizer, |
| BitsAndBytesConfig, |
| Trainer, |
| TrainingArguments, |
| set_seed, |
| ) |
|
|
| SYSTEM_PREFIX = ( |
| "You are GravityLLM, a Spatial9 scene generation model. " |
| "Given music constraints and stem features, output ONLY valid Spatial9Scene JSON. " |
| "Do not return markdown. Do not explain your answer. " |
| "Respect hard constraints such as object budgets, anchor positions, and low-end centering.\n\n" |
| ) |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description="Fine-tune GravityLLM for Spatial9 scene generation.") |
| parser.add_argument("--model", type=str, default="Qwen/Qwen2.5-1.5B-Instruct") |
| parser.add_argument("--train_file", type=str, default="data/train.jsonl") |
| parser.add_argument("--valid_file", type=str, default="data/valid.jsonl") |
| parser.add_argument("--output_dir", type=str, default="outputs/GravityLLM-Qwen2.5-1.5B-S9") |
| parser.add_argument("--max_length", type=int, default=2048) |
|
|
| parser.add_argument("--num_train_epochs", type=float, default=1.0) |
| parser.add_argument("--learning_rate", type=float, default=2e-4) |
| parser.add_argument("--train_batch_size", type=int, default=1) |
| parser.add_argument("--eval_batch_size", type=int, default=1) |
| parser.add_argument("--gradient_accumulation_steps", type=int, default=16) |
| parser.add_argument("--warmup_ratio", type=float, default=0.03) |
| parser.add_argument("--weight_decay", type=float, default=0.0) |
| parser.add_argument("--logging_steps", type=int, default=10) |
| parser.add_argument("--save_steps", type=int, default=200) |
| parser.add_argument("--eval_steps", type=int, default=200) |
| parser.add_argument("--seed", type=int, default=42) |
|
|
| parser.add_argument("--lora", action="store_true", help="Enable LoRA adapters.") |
| parser.add_argument("--qlora", action="store_true", help="Enable 4-bit QLoRA training.") |
| parser.add_argument("--lora_r", type=int, default=16) |
| parser.add_argument("--lora_alpha", type=int, default=32) |
| parser.add_argument("--lora_dropout", type=float, default=0.05) |
|
|
| parser.add_argument("--bf16", action="store_true") |
| parser.add_argument("--fp16", action="store_true") |
|
|
| parser.add_argument("--push_to_hub", action="store_true") |
| parser.add_argument("--hub_model_id", type=str, default=None) |
| parser.add_argument("--hub_private_repo", action="store_true") |
| return parser.parse_args() |
|
|
|
|
| def load_jsonl(file_path: str): |
| return load_dataset("json", data_files=file_path, split="train") |
|
|
|
|
| def format_prompt(raw_prompt: str) -> str: |
| raw_prompt = raw_prompt.strip() |
| if raw_prompt.lower().startswith("gravityllm:"): |
| raw_prompt = raw_prompt.split(":", 1)[1].strip() |
| return SYSTEM_PREFIX + raw_prompt + "\n\nOUTPUT:\n" |
|
|
|
|
| def tokenize_example(example: Dict[str, str], tokenizer, max_length: int) -> Dict[str, List[int]]: |
| prompt_text = format_prompt(example["prompt"]) |
| completion_text = example["completion"].strip() |
|
|
| prompt_ids = tokenizer(prompt_text, add_special_tokens=False)["input_ids"] |
| completion_ids = tokenizer(completion_text + tokenizer.eos_token, add_special_tokens=False)["input_ids"] |
|
|
| input_ids = prompt_ids + completion_ids |
| labels = [-100] * len(prompt_ids) + completion_ids |
|
|
| if len(input_ids) > max_length: |
| input_ids = input_ids[:max_length] |
| labels = labels[:max_length] |
|
|
| attention_mask = [1] * len(input_ids) |
| return { |
| "input_ids": input_ids, |
| "attention_mask": attention_mask, |
| "labels": labels, |
| } |
|
|
|
|
| @dataclass |
| class CausalDataCollator: |
| pad_token_id: int |
| label_pad_token_id: int = -100 |
|
|
| def __call__(self, features): |
| max_len = max(len(f["input_ids"]) for f in features) |
|
|
| input_ids = [] |
| attention_mask = [] |
| labels = [] |
|
|
| for f in features: |
| pad_len = max_len - len(f["input_ids"]) |
| input_ids.append(f["input_ids"] + [self.pad_token_id] * pad_len) |
| attention_mask.append(f["attention_mask"] + [0] * pad_len) |
| labels.append(f["labels"] + [self.label_pad_token_id] * pad_len) |
|
|
| batch = { |
| "input_ids": torch.tensor(input_ids, dtype=torch.long), |
| "attention_mask": torch.tensor(attention_mask, dtype=torch.long), |
| "labels": torch.tensor(labels, dtype=torch.long), |
| } |
| return batch |
|
|
|
|
| def prepare_model(args: argparse.Namespace): |
| model_kwargs = {} |
| if args.qlora: |
| compute_dtype = torch.bfloat16 if args.bf16 else torch.float16 |
| model_kwargs["quantization_config"] = BitsAndBytesConfig( |
| load_in_4bit=True, |
| bnb_4bit_quant_type="nf4", |
| bnb_4bit_use_double_quant=True, |
| bnb_4bit_compute_dtype=compute_dtype, |
| ) |
| model_kwargs["device_map"] = "auto" |
|
|
| model = AutoModelForCausalLM.from_pretrained( |
| args.model, |
| torch_dtype=torch.bfloat16 if args.bf16 else (torch.float16 if args.fp16 else None), |
| trust_remote_code=True, |
| **model_kwargs, |
| ) |
| model.config.use_cache = False |
|
|
| if args.qlora: |
| model = prepare_model_for_kbit_training(model) |
|
|
| if args.lora or args.qlora: |
| lora_config = LoraConfig( |
| r=args.lora_r, |
| lora_alpha=args.lora_alpha, |
| lora_dropout=args.lora_dropout, |
| bias="none", |
| task_type="CAUSAL_LM", |
| target_modules="all-linear", |
| ) |
| model = get_peft_model(model, lora_config) |
| model.print_trainable_parameters() |
|
|
| return model |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| os.makedirs(args.output_dir, exist_ok=True) |
| set_seed(args.seed) |
|
|
| tokenizer = AutoTokenizer.from_pretrained(args.model, use_fast=True, trust_remote_code=True) |
| tokenizer.padding_side = "right" |
| if tokenizer.pad_token is None: |
| tokenizer.pad_token = tokenizer.eos_token |
|
|
| train_ds = load_jsonl(args.train_file) |
| valid_ds = load_jsonl(args.valid_file) if args.valid_file and Path(args.valid_file).exists() else None |
|
|
| train_ds = train_ds.map( |
| lambda row: tokenize_example(row, tokenizer, args.max_length), |
| remove_columns=train_ds.column_names, |
| desc="Tokenizing train set", |
| ) |
| if valid_ds is not None: |
| valid_ds = valid_ds.map( |
| lambda row: tokenize_example(row, tokenizer, args.max_length), |
| remove_columns=valid_ds.column_names, |
| desc="Tokenizing valid set", |
| ) |
|
|
| model = prepare_model(args) |
|
|
| training_args = TrainingArguments( |
| output_dir=args.output_dir, |
| overwrite_output_dir=True, |
| num_train_epochs=args.num_train_epochs, |
| learning_rate=args.learning_rate, |
| per_device_train_batch_size=args.train_batch_size, |
| per_device_eval_batch_size=args.eval_batch_size, |
| gradient_accumulation_steps=args.gradient_accumulation_steps, |
| warmup_ratio=args.warmup_ratio, |
| weight_decay=args.weight_decay, |
| logging_steps=args.logging_steps, |
| save_steps=args.save_steps, |
| eval_steps=args.eval_steps, |
| evaluation_strategy="steps" if valid_ds is not None else "no", |
| save_strategy="steps", |
| bf16=args.bf16, |
| fp16=args.fp16, |
| report_to="none", |
| gradient_checkpointing=True, |
| lr_scheduler_type="cosine", |
| optim="paged_adamw_32bit" if (args.lora or args.qlora) else "adamw_torch", |
| max_grad_norm=1.0, |
| push_to_hub=args.push_to_hub, |
| hub_model_id=args.hub_model_id, |
| hub_private_repo=args.hub_private_repo, |
| hub_strategy="end" if args.push_to_hub else "every_save", |
| ) |
|
|
| trainer = Trainer( |
| model=model, |
| args=training_args, |
| train_dataset=train_ds, |
| eval_dataset=valid_ds, |
| data_collator=CausalDataCollator(pad_token_id=tokenizer.pad_token_id), |
| tokenizer=tokenizer, |
| ) |
|
|
| train_result = trainer.train() |
| trainer.save_model(args.output_dir) |
| tokenizer.save_pretrained(args.output_dir) |
|
|
| metrics = train_result.metrics |
| with open(Path(args.output_dir) / "training_metrics.json", "w", encoding="utf-8") as f: |
| json.dump(metrics, f, indent=2) |
|
|
| run_meta = vars(args).copy() |
| run_meta["train_examples"] = len(train_ds) |
| run_meta["valid_examples"] = len(valid_ds) if valid_ds is not None else 0 |
| with open(Path(args.output_dir) / "run_config.json", "w", encoding="utf-8") as f: |
| json.dump(run_meta, f, indent=2) |
|
|
| if args.push_to_hub: |
| trainer.push_to_hub(commit_message="Add GravityLLM fine-tuned adapter") |
| print(f"Training complete. Artifacts saved to: {args.output_dir}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
