Update training script for synthetic dataset

train.py

@@ -1,126 +1,115 @@
 """
 Parametric Floorplan Generation Model Training
-Based on: DStruct2Design (arXiv:2407.15723)
-Approach: Fine-tune an instruction-tuned LLM (Qwen2.5-1.5B-Instruct) with LoRA
-          to generate JSON floorplan structures from parametric constraint prompts.
-Dataset: ludolara/DStruct2Design (10k train, 1k val, 1k test)
+Based on: DStruct2Design (arXiv:2407.15723) approach
+Dataset: Custom synthetic dataset generated to match user's ProjectCreate schema
 """
 import os
 import json
 import torch
-from datasets import load_dataset, DatasetDict
+from datasets import load_dataset, load_from_disk, DatasetDict
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import LoraConfig, TaskType
 from trl import SFTTrainer, SFTConfig
 
-def format_constraints(example):
-    room_count = example["room_count"]
-    total_area = example["total_area"]
-    room_types = example["room_types"]
-    edges = example.get("edges", [])
-    rooms = example.get("rooms", [])
-    lines = [
-        f"Generate a floor plan with {room_count} rooms and a total area of {total_area} square meters.",
-        f"The room types are: {', '.join(room_types)}."
-    ]
-    if rooms:
-        lines.append("Room details:")
-        for i, room in enumerate(rooms):
-            lines.append(f"  - Room {i+1} ({room.get('room_type','unknown')}): area ~{room.get('area','unspecified')} m², width ~{room.get('width','unspecified')} m, height ~{room.get('height','unspecified')} m")
-    if edges:
-        lines.append(f"Adjacency requirements (room indices): {edges}")
-    return "\n".join(lines)
+# -----------------------------------------------------------------------------
+# Configuration
+# -----------------------------------------------------------------------------
+MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
+OUTPUT_DIR = "/app/floorplan-model"
+HUB_MODEL_ID = os.environ.get("HF_TRAINER_HUB_MODEL_ID", "Karthik8nitt/parametric-floorplan-generator")
+DATASET_PATH = os.environ.get("DATASET_PATH", "/app/floorplan_synthetic_dataset")
 
-def format_floorplan_output(example):
-    return json.dumps({
-        "rooms": [{"room_type": r["room_type"], "area": r["area"], "width": r["width"],
-                   "height": r["height"], "floor_polygon": r["floor_polygon"],
-                   "is_regular": r.get("is_regular", 0)} for r in example["rooms"]],
-        "edges": example.get("edges", []),
-        "room_count": example["room_count"],
-        "total_area": example["total_area"],
-        "room_types": example["room_types"],
-    }, indent=2)
+# -----------------------------------------------------------------------------
+# Load data
+# -----------------------------------------------------------------------------
+print("Loading dataset...")
+if os.path.exists(DATASET_PATH):
+    dataset = load_from_disk(DATASET_PATH)
+else:
+    # Fallback: load from HF if pre-uploaded
+    dataset = load_dataset("Karthik8nitt/floorplan-synthetic-dataset")
 
-def convert(example):
-    return {"prompt": format_constraints(example), "completion": format_floorplan_output(example)}
+print(f"Train: {len(dataset['train'])}, Val: {len(dataset['validation'])}, Test: {len(dataset['test'])}")
 
-def main():
-    model_id = "Qwen/Qwen2.5-1.5B-Instruct"
-    hub_model_id = os.environ.get("HF_TRAINER_HUB_MODEL_ID", "Karthik8nitt/parametric-floorplan-generator")
-    output_dir = "/app/floorplan-model"
+# -----------------------------------------------------------------------------
+# Load tokenizer & model
+# -----------------------------------------------------------------------------
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
 
-    print("Loading DStruct2Design dataset...")
-    dataset = load_dataset("ludolara/DStruct2Design")
-    processed = {split: dataset[split].map(convert, remove_columns=dataset[split].column_names) for split in dataset.keys()}
+print("Loading model...")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    trust_remote_code=True,
+)
 
-    print("Loading tokenizer...")
-    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
+# -----------------------------------------------------------------------------
+# LoRA config
+# -----------------------------------------------------------------------------
+peft_config = LoraConfig(
+    r=16,
+    lora_alpha=32,
+    lora_dropout=0.05,
+    bias="none",
+    task_type=TaskType.CAUSAL_LM,
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+)
 
-    print("Loading model...")
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-        trust_remote_code=True,
-    )
+# -----------------------------------------------------------------------------
+# Training arguments
+# -----------------------------------------------------------------------------
+training_args = SFTConfig(
+    output_dir=OUTPUT_DIR,
+    num_train_epochs=5,
+    per_device_train_batch_size=4,
+    per_device_eval_batch_size=4,
+    gradient_accumulation_steps=4,
+    learning_rate=1e-4,
+    lr_scheduler_type="cosine",
+    warmup_ratio=0.1,
+    logging_steps=10,
+    eval_strategy="steps",
+    eval_steps=100,
+    save_strategy="steps",
+    save_steps=100,
+    save_total_limit=3,
+    max_seq_length=4096,
+    bf16=True,
+    gradient_checkpointing=True,
+    report_to="trackio",
+    run_name="floorplan-qwen1.5b-lora",
+    project="parametric-floorplan",
+    hub_model_id=HUB_MODEL_ID,
+    push_to_hub=True,
+    completion_only_loss=True,
+    disable_tqdm=True,
+    logging_first_step=True,
+    seed=42,
+)
 
-    peft_config = LoraConfig(
-        r=16,
-        lora_alpha=32,
-        lora_dropout=0.05,
-        bias="none",
-        task_type=TaskType.CAUSAL_LM,
-        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
-    )
+# -----------------------------------------------------------------------------
+# Trainer
+# -----------------------------------------------------------------------------
+trainer = SFTTrainer(
+    model=model,
+    args=training_args,
+    train_dataset=dataset["train"],
+    eval_dataset=dataset["validation"],
+    peft_config=peft_config,
+    processing_class=tokenizer,
+)
 
-    training_args = SFTConfig(
-        output_dir=output_dir,
-        num_train_epochs=5,
-        per_device_train_batch_size=4,
-        per_device_eval_batch_size=4,
-        gradient_accumulation_steps=4,
-        learning_rate=1e-4,
-        lr_scheduler_type="cosine",
-        warmup_ratio=0.1,
-        logging_steps=10,
-        eval_strategy="steps",
-        eval_steps=100,
-        save_strategy="steps",
-        save_steps=100,
-        save_total_limit=3,
-        max_seq_length=2048,
-        bf16=True,
-        gradient_checkpointing=True,
-        report_to="trackio",
-        run_name="floorplan-qwen1.5b-lora",
-        project="parametric-floorplan",
-        hub_model_id=hub_model_id,
-        push_to_hub=True,
-        completion_only_loss=True,
-        disable_tqdm=True,
-        logging_first_step=True,
-        seed=42,
-    )
+# -----------------------------------------------------------------------------
+# Train
+# -----------------------------------------------------------------------------
+print("Starting training...")
+trainer.train()
 
-    trainer = SFTTrainer(
-        model=model,
-        args=training_args,
-        train_dataset=processed["train"],
-        eval_dataset=processed["validation"],
-        peft_config=peft_config,
-        processing_class=tokenizer,
-    )
-
-    print("Starting training...")
-    trainer.train()
-
-    print("Saving and pushing model...")
-    trainer.save_model(os.path.join(output_dir, "final"))
-    trainer.push_to_hub()
-    print(f"Done! Model at https://huggingface.co/{hub_model_id}")
-
-if __name__ == "__main__":
-    main()
+print("Saving and pushing model...")
+trainer.save_model(os.path.join(OUTPUT_DIR, "final"))
+trainer.push_to_hub()
+print(f"Done! Model at https://huggingface.co/{HUB_MODEL_ID}")