Add training script

train.py

@@ -0,0 +1,126 @@
+"""
+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)
+"""
+import os
+import json
+import torch
+from datasets import load_dataset, 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)
+
+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)
+
+def convert(example):
+    return {"prompt": format_constraints(example), "completion": format_floorplan_output(example)}
+
+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"
+
+    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 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 model...")
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+        trust_remote_code=True,
+    )
+
+    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"],
+    )
+
+    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,
+    )
+
+    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()