Add Modal training script

modal_train.py

@@ -0,0 +1,429 @@
+"""
+Train the parametric floorplan model on Modal.com GPU instances.
+Usage: modal deploy modal_train.py
+       modal run modal_train.py
+"""
+import modal
+import os
+
+# ---------------------------------------------------------------------------
+# Modal Image & App Setup
+# ---------------------------------------------------------------------------
+image = (
+    modal.Image.debian_slim(python_version="3.10")
+    .pip_install(
+        "transformers>=4.45.0",
+        "trl>=0.15.0",
+        "torch>=2.3.0",
+        "datasets>=2.20.0",
+        "peft>=0.12.0",
+        "accelerate>=0.33.0",
+        "trackio>=0.3.0",
+        "huggingface_hub",
+    )
+    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
+)
+
+app = modal.App("floorplan-trainer", image=image)
+
+# ---------------------------------------------------------------------------
+# Persistent volume for dataset & model checkpoints
+# ---------------------------------------------------------------------------
+vol = modal.Volume.from_name("floorplan-data", create_if_missing=True)
+MODEL_VOL = modal.Volume.from_name("floorplan-model", create_if_missing=True)
+
+# ---------------------------------------------------------------------------
+# Secrets
+# ---------------------------------------------------------------------------
+hf_secret = modal.Secret.from_name("huggingface-token")
+
+# ---------------------------------------------------------------------------
+# Dataset Generation (CPU)
+# ---------------------------------------------------------------------------
+@app.function(volumes={"/data": vol}, secrets=[hf_secret], timeout=1800)
+def generate_dataset_on_modal():
+    """Generate the synthetic dataset inside Modal volume."""
+    import json, random
+    from typing import List, Dict, Any
+    from datasets import Dataset, DatasetDict
+
+    DIRECTIONS = ["N", "S", "E", "W"]
+    CITIES = ["Delhi", "Mumbai", "Bangalore", "Chennai", "Hyderabad", "Pune", "Kolkata", "Ahmedabad", "Jaipur", "other"]
+    MUNICIPALITIES = ["MC", "MDA", "PMA", "BDA", "GHMC", "BBMP", "MCD", "KMC", "JDA", None]
+
+    def rect_polygon(x, y, w, d):
+        return [[round(v,2) for v in p] for p in [[x,y],[x+w,y],[x+w,y+d],[x,y+d]]]
+    def polygon_area(poly):
+        n = len(poly); area = 0.0
+        for i in range(n):
+            x1,y1 = poly[i]; x2,y2 = poly[(i+1)%n]
+            area += x1*y2 - x2*y1
+        return abs(area)/2.0
+    def polygon_bbox(poly):
+        xs = [p[0] for p in poly]; ys = [p[1] for p in poly]
+        return min(xs), min(ys), max(xs), max(ys)
+
+    def make_plot_boundary(params):
+        shape = params.get("plot_shape","rectangular")
+        L,W = params["plot_length"], params["plot_width"]
+        if shape == "rectangular": return rect_polygon(0,0,L,W)
+        if shape == "l_shaped":
+            cw = params.get("cutout_width",L*0.3); ch = params.get("cutout_height",W*0.3)
+            corner = params.get("cutout_corner","NE")
+            if corner == "NE": return [[0,0],[L,0],[L,W-ch],[L-cw,W-ch],[L-cw,W],[0,W]]
+            if corner == "NW": return [[0,0],[L,0],[L,W],[cw,W],[cw,W-ch],[0,W-ch]]
+            if corner == "SE": return [[0,0],[L-cw,0],[L-cw,ch],[L,ch],[L,W],[0,W]]
+            if corner == "SW": return [[cw,0],[L,0],[L,W],[0,W],[0,ch],[cw,ch]]
+            return rect_polygon(0,0,L,W)
+        if shape == "trapezoid":
+            fw = params.get("plot_front_width",L); rw = params.get("plot_rear_width",L*0.8)
+            off = params.get("plot_side_offset",0.0)
+            return [[0,0],[fw,0],[fw-off,W],[-off,W]]
+        return rect_polygon(0,0,L,W)
+
+    def make_buildable_boundary(plot_poly, params):
+        sf = params.get("setback_front",1.5); sr = params.get("setback_rear",1.0)
+        sl = params.get("setback_left",1.0); srt = params.get("setback_right",1.0)
+        minx,miny,maxx,maxy = polygon_bbox(plot_poly)
+        return rect_polygon(minx+sl,miny+sf,maxx-minx-sl-srt,maxy-miny-sf-sr)
+
+    def distribute_width(total,n):
+        base = total/n; bays = []; remaining = total
+        for i in range(n):
+            if i == n-1: bays.append(round(remaining,2))
+            else:
+                bay = max(2.4, round(base,1))
+                bay = min(bay, remaining - 2.4*(n-i-1))
+                bays.append(bay); remaining -= bay
+        return bays
+
+    def generate_room_specs(params):
+        specs = []
+        num_bed = params["num_bedrooms"]; num_toi = params["toilets"]
+        has_pooja = params.get("has_pooja",False); has_study = params.get("has_study",False)
+        has_balc = params.get("has_balcony",False); parking = params.get("parking",False)
+        num_floors = params.get("num_floors",1); has_stilt = params.get("has_stilt",False)
+        has_basement = params.get("has_basement",False); custom = params.get("custom_room_config",None) or []
+        def add(rtype,rid,name,zone,floor="gf",target_area=None):
+            specs.append({"id":rid,"type":rtype,"name":name,"zone":zone,"floor":floor,"target_area":target_area})
+        add("living","living_1","Living Room","front","gf",18)
+        add("kitchen","kitchen_1","Kitchen","mid","gf",9)
+        add("dining","dining_1","Dining Area","front","gf",12)
+        add("toilet","toilet_common","Common Toilet","back","gf",3.5)
+        for i in range(num_bed):
+            floor = "gf" if i < num_bed-(num_floors-1) else ("ff" if num_floors>1 else "gf")
+            if i==0: add("master_bedroom",f"bedroom_{i+1}","Master Bedroom","back",floor,16)
+            else: add("bedroom",f"bedroom_{i+1}",f"Bedroom {i+1}","back",floor,12)
+        for i in range(num_toi):
+            floor = "gf" if i < num_toi-(num_floors-1) else ("ff" if num_floors>1 else "gf")
+            add("toilet",f"toilet_{i+1}",f"Toilet {i+1}","back",floor,3.5)
+        if has_pooja: add("pooja","pooja_1","Pooja Room","back","gf",2.5)
+        if has_study: floor = "ff" if num_floors>1 else "gf"; add("study","study_1","Study Room","back",floor,8)
+        if has_balc: add("balcony","balcony_1","Balcony","side","gf",5)
+        add("staircase","stairs_1","Staircase","mid","gf",10)
+        if has_stilt or parking: add("parking","parking_1","Parking","side","stilt",15); add("staircase","stairs_stilt","Staircase (Stilt)","mid","stilt",10)
+        if has_basement: add("store","store_base","Storage","back","basement",8); add("staircase","stairs_base","Staircase (Basement)","mid","basement",10)
+        if num_floors >= 2:
+            ff_beds = max(0,num_bed-1)
+            for i in range(ff_beds): add("bedroom",f"bedroom_ff_{i+1}",f"Bedroom {num_bed-ff_beds+i+1}","back","ff",12)
+            add("living","living_ff","Family Lounge","front","ff",14)
+            add("toilet","toilet_ff","Common Toilet (FF)","back","ff",3.5)
+            if has_study and num_floors>=2: add("study","study_ff","Study Room","back","ff",8)
+            if has_balc: add("balcony","balcony_ff","Balcony (FF)","side","ff",5)
+        if num_floors >= 3:
+            sf_beds = max(0,num_bed-2)
+            for i in range(sf_beds): add("bedroom",f"bedroom_sf_{i+1}",f"Bedroom {num_bed-sf_beds-ff_beds+i+1}","back","sf",12)
+            add("living","living_sf","Terrace Lounge","front","sf",12)
+            add("toilet","toilet_sf","Common Toilet (SF)","back","sf",3.5)
+            if has_balc: add("balcony","balcony_sf","Balcony (SF)","side","sf",5)
+        for i,cr in enumerate(custom):
+            rtype = cr.get("type","room").lower().replace(" ","_")
+            floor_map = {"basement":"basement","stilt":"stilt","gf":"gf","ff":"ff","sf":"sf","either":"gf"}
+            floor = floor_map.get(cr.get("floor_preference","either"),"gf")
+            add(rtype,f"custom_{i+1}",cr.get("name",f"Custom Room {i+1}"),"mid",floor,cr.get("min_area_sqm",10))
+        return specs
+
+    def place_rooms(buildable_poly, rooms_spec, vastu, road_side, north_dir):
+        minx,miny,maxx,maxy = polygon_bbox(buildable_poly)
+        bw,bd = maxx-minx,maxy-miny
+        num_bays = max(2,min(4,int(bw/3.0)))
+        bay_widths = distribute_width(bw, num_bays)
+        def place_row(room_list,row_depth,y_start):
+            x_cursor = miny
+            placed_in_row = []
+            for i,room in enumerate(room_list):
+                if i >= len(bay_widths): break
+                w = bay_widths[i]; d = row_depth
+                target = room.get("target_area",w*d)
+                if target>0 and w>0:
+                    adj_d = min(max(target/w,2.4),row_depth)
+                    d = round(adj_d,2)
+                poly = rect_polygon(round(x_cursor,2),round(y_start,2),round(w,2),round(d,2))
+                area = polygon_area(poly)
+                placed_in_row.append({
+                    "id":room["id"],"type":room["type"],"name":room["name"],"floor":room.get("floor","gf"),
+                    "polygon":poly,"area_sqm":round(area,2),
+                    "dimensions":{"width":round(w,2),"depth":round(d,2)},
+                    "position":{"x":round(x_cursor+w/2,2),"y":round(y_start+d/2,2)},
+                })
+                x_cursor += w
+            return placed_in_row
+        all_rows = []
+        front_types = [r for r in rooms_spec if r["type"] in ("living","dining")]
+        if front_types: all_rows.append((front_types, bd*0.35))
+        mid_types = [r for r in rooms_spec if r["type"] in ("kitchen","utility","staircase","corridor","store")]
+        if mid_types: all_rows.append((mid_types, bd*0.3))
+        back_types = [r for r in rooms_spec if r["type"] in ("bedroom","master_bedroom","toilet","pooja","study")]
+        if back_types: all_rows.append((back_types, bd*0.35))
+        y_cursor = miny
+        for room_list,row_depth in all_rows:
+            yield from place_row(room_list, row_depth, y_cursor)
+            y_cursor += row_depth
+        side_types = [r for r in rooms_spec if r["type"] in ("balcony","parking")]
+        for room in side_types:
+            if room["type"] == "balcony" and list(all_rows):
+                poly = rect_polygon(minx, maxy, min(bw/num_bays, 2.0), min(bd*0.15, 1.5))
+                yield {
+                    "id": room["id"], "type": "balcony", "name": room.get("name","Balcony"), "floor": room.get("floor","gf"),
+                    "polygon": [[round(v,2) for v in p] for p in poly], "area_sqm": round(polygon_area(poly),2),
+                    "dimensions": {"width": round(min(bw/num_bays,2.0),2), "depth": round(min(bd*0.15,1.5),2)},
+                    "position": {"x": round(minx+min(bw/num_bays,2.0)/2,2), "y": round(maxy+min(bd*0.15,1.5)/2,2)},
+                }
+            elif room["type"] == "parking":
+                poly = rect_polygon(minx, miny, min(bw,7.5), min(bd*0.25,6.0))
+                yield {
+                    "id": room["id"], "type": "parking", "name": room.get("name","Parking"), "floor": room.get("floor","stilt"),
+                    "polygon": [[round(v,2) for v in p] for p in poly], "area_sqm": round(polygon_area(poly),2),
+                    "dimensions": {"width": round(min(bw,7.5),2), "depth": round(min(bd*0.25,6.0),2)},
+                    "position": {"x": round(minx+min(bw,7.5)/2,2), "y": round(miny+min(bd*0.25,6.0)/2,2)},
+                }
+
+    def generate_openings(rooms, road_side):
+        doors,windows = [],[]
+        entrance = [r for r in rooms if r["type"]=="living" and r["floor"]=="gf"]
+        if entrance:
+            lr = entrance[0]; mx,my,Mx,My = polygon_bbox(lr["polygon"])
+            if road_side in ("N","S"):
+                x = round((mx+Mx)/2 - 0.45, 2); y = My if road_side=="N" else my
+                doors.append({"id":"door_main","type":"main_entrance","width":0.9,"from":"outside","to":lr["id"],"position":[x,y],"orientation":"horizontal"})
+            else:
+                x = Mx if road_side=="E" else mx; y = round((my+My)/2 - 0.45, 2)
+                doors.append({"id":"door_main","type":"main_entrance","width":0.9,"from":"outside","to":lr["id"],"position":[x,y],"orientation":"vertical"})
+        for i,r1 in enumerate(rooms):
+            for r2 in rooms[i+1:]:
+                if r1["floor"] != r2["floor"]: continue
+                m1x,m1y,M1x,M1y = polygon_bbox(r1["polygon"])
+                m2x,m2y,M2x,M2y = polygon_bbox(r2["polygon"])
+                share_x = not (M1x < m2x or M2x < m1x); share_y = not (M1y < m2y or M2y < m1y)
+                if share_x and abs(M1y-m2y) < 0.3:
+                    x = round(max(m1x,m2x)+0.3,2); y = round(M1y,2)
+                    doors.append({"id":f"door_{r1['id']}_{r2['id']}","type":"internal","width":0.75,"from":r1["id"],"to":r2["id"],"position":[x,y],"orientation":"horizontal"})
+                elif share_y and abs(M1x-m2x) < 0.3:
+                    x = round(M1x,2); y = round(max(m1y,m2y)+0.3,2)
+                    doors.append({"id":f"door_{r1['id']}_{r2['id']}","type":"internal","width":0.75,"from":r1["id"],"to":r2["id"],"position":[x,y],"orientation":"vertical"})
+        for r in rooms:
+            if r["type"] in ("living","bedroom","master_bedroom","dining","kitchen","study"):
+                mx,my,Mx,My = polygon_bbox(r["polygon"])
+                if (Mx-mx) >= (My-my):
+                    y = round((my+My)/2,2); cx = (mx+Mx)/2; x = mx if abs(mx-cx) > abs(Mx-cx) else Mx
+                    windows.append({"id":f"win_{r['id']}","room":r["id"],"width":1.2,"height":1.5,"position":[round(x,2),y],"orientation":"vertical"})
+                else:
+                    x = round((mx+Mx)/2,2); cy = (my+My)/2; y = my if abs(my-cy) > abs(My-cy) else My
+                    windows.append({"id":f"win_{r['id']}","room":r["id"],"width":1.5,"height":1.2,"position":[x,round(y,2)],"orientation":"horizontal"})
+        return doors,windows
+
+    def generate_example(seed=None):
+        if seed is not None: random.seed(seed)
+        plot_length = round(random.uniform(8.0,25.0),1)
+        plot_width = round(random.uniform(7.0,20.0),1)
+        setback_front = round(random.uniform(1.0,3.0),1); setback_rear = round(random.uniform(0.5,2.0),1)
+        setback_left = round(random.uniform(0.5,2.0),1); setback_right = round(random.uniform(0.5,2.0),1)
+        road_side = random.choice(DIRECTIONS); north_direction = random.choice(DIRECTIONS)
+        num_bedrooms = random.randint(1,4); toilets = random.randint(1,num_bedrooms+1)
+        parking = random.choice([True,False]); city = random.choice(CITIES)
+        vastu_enabled = random.choice([True,False]); road_width_m = round(random.uniform(6.0,18.0),1)
+        has_pooja = random.choice([True,False]); has_study = random.choice([True,False])
+        has_balcony = random.choice([True,False])
+        plot_shape = random.choice(["rectangular"]*8 + ["l_shaped"]*1 + ["trapezoid"]*1)
+        plot_front_width = plot_length if plot_shape!="trapezoid" else round(plot_length*random.uniform(0.8,1.0),1)
+        plot_rear_width = plot_length if plot_shape!="trapezoid" else round(plot_length*random.uniform(0.7,1.0),1)
+        plot_side_offset = 0.0 if plot_shape!="trapezoid" else round(random.uniform(-1.0,1.0),1)
+        cutout_corner = random.choice(["NE","NW","SE","SW"])
+        cutout_width = round(plot_length*random.uniform(0.15,0.35),1) if plot_shape=="l_shaped" else 0.0
+        cutout_height = round(plot_width*random.uniform(0.15,0.35),1) if plot_shape=="l_shaped" else 0.0
+        num_floors = random.choices([1,2,3], weights=[5,3,1])[0]
+        has_stilt = random.choice([True,False]) if num_floors>1 else False
+        has_basement = random.choice([True,False]); municipality = random.choice(MUNICIPALITIES)
+        min_buildable = 5.0
+        if plot_length - setback_front - setback_rear < min_buildable:
+            setback_front = min(setback_front,(plot_length-min_buildable)/2)
+            setback_rear = min(setback_rear,(plot_length-min_buildable)/2)
+        if plot_width - setback_left - setback_right < min_buildable:
+            setback_left = min(setback_left,(plot_width-min_buildable)/2)
+            setback_right = min(setback_right,(plot_width-min_buildable)/2)
+        params = {
+            "name": f"Project_{random.randint(1000,9999)}",
+            "plot_length": plot_length, "plot_width": plot_width,
+            "setback_front": round(setback_front,1), "setback_rear": round(setback_rear,1),
+            "setback_left": round(setback_left,1), "setback_right": round(setback_right,1),
+            "road_side": road_side, "north_direction": north_direction,
+            "num_bedrooms": num_bedrooms, "toilets": toilets, "parking": parking,
+            "city": city, "vastu_enabled": vastu_enabled, "road_width_m": road_width_m,
+            "has_pooja": has_pooja, "has_study": has_study, "has_balcony": has_balcony,
+            "plot_shape": plot_shape,
+            "plot_front_width": plot_front_width if plot_shape=="trapezoid" else None,
+            "plot_rear_width": plot_rear_width if plot_shape=="trapezoid" else None,
+            "plot_side_offset": plot_side_offset if plot_shape=="trapezoid" else None,
+            "plot_corners": None,
+            "cutout_corner": cutout_corner, "cutout_width": cutout_width, "cutout_height": cutout_height,
+            "num_floors": num_floors, "has_stilt": has_stilt, "has_basement": has_basement,
+            "municipality": municipality, "custom_room_config": None, "team_id": None,
+        }
+        plot_boundary = make_plot_boundary(params)
+        buildable_boundary = make_buildable_boundary(plot_boundary, params)
+        room_specs = generate_room_specs(params)
+        rooms = list(place_rooms(buildable_boundary, room_specs, vastu_enabled, road_side, north_direction))
+        doors, windows = generate_openings(rooms, road_side)
+        total_area = sum(r["area_sqm"] for r in rooms if r["floor"] not in ("stilt","basement"))
+        built_up_area = sum(r["area_sqm"] for r in rooms)
+        floorplan = {
+            "project_name": params["name"],
+            "plot": {
+                "shape": plot_shape, "outer_boundary": plot_boundary,
+                "setbacks": {"front":params["setback_front"],"rear":params["setback_rear"],"left":params["setback_left"],"right":params["setback_right"]},
+                "buildable_boundary": buildable_boundary,
+                "road_side": road_side, "north_direction": north_direction,
+                "plot_length": plot_length, "plot_width": plot_width,
+            },
+            "rooms": rooms, "doors": doors, "windows": windows,
+            "dimensions": {
+                "total_built_up_area_sqm": round(built_up_area,2),
+                "total_carpet_area_sqm": round(total_area,2),
+                "ground_floor_area_sqm": round(sum(r["area_sqm"] for r in rooms if r["floor"]=="gf"),2),
+                "first_floor_area_sqm": round(sum(r["area_sqm"] for r in rooms if r["floor"]=="ff"),2),
+                "second_floor_area_sqm": round(sum(r["area_sqm"] for r in rooms if r["floor"]=="sf"),2),
+                "stilt_area_sqm": round(sum(r["area_sqm"] for r in rooms if r["floor"]=="stilt"),2),
+                "basement_area_sqm": round(sum(r["area_sqm"] for r in rooms if r["floor"]=="basement"),2),
+            },
+            "meta": {"num_floors": num_floors, "has_stilt": has_stilt, "has_basement": has_basement,
+                     "vastu_enabled": vastu_enabled, "city": city, "municipality": municipality},
+        }
+        lines = [
+            f"Generate a floor plan for project '{params['name']}'.",
+            f"Plot dimensions: {plot_length}m x {plot_width}m, shape: {plot_shape}.",
+            f"Setbacks: front={params['setback_front']}m, rear={params['setback_rear']}m, left={params['setback_left']}m, right={params['setback_right']}m.",
+            f"Road side: {road_side}, North direction: {north_direction}.",
+            f"Requirements: {num_bedrooms} bedrooms, {toilets} toilets.",
+        ]
+        if parking: lines.append("Parking is required.")
+        if has_pooja: lines.append("Include a Pooja room.")
+        if has_study: lines.append("Include a Study room.")
+        if has_balcony: lines.append("Include a Balcony.")
+        if has_stilt: lines.append("Stilt parking required.")
+        if has_basement: lines.append("Include a basement.")
+        lines.append(f"Number of floors: {num_floors} (1=G, 2=G+1, 3=G+2).")
+        if vastu_enabled: lines.append("Vastu compliance is enabled.")
+        lines.append(f"City: {city}, Municipality: {municipality or 'N/A'}.")
+        prompt = "\n".join(lines)
+        return {"prompt": prompt, "completion": json.dumps(floorplan, indent=2), "params": params}
+
+    ds = DatasetDict({
+        "train": Dataset.from_list([generate_example(seed=i) for i in range(5000)]),
+        "validation": Dataset.from_list([generate_example(seed=100000+i) for i in range(500)]),
+        "test": Dataset.from_list([generate_example(seed=200000+i) for i in range(500)]),
+    })
+    ds.save_to_disk("/data/floorplan_synthetic_dataset")
+    print("Dataset saved to /data/floorplan_synthetic_dataset")
+    print(f"Train: {len(ds['train'])}, Val: {len(ds['validation'])}, Test: {len(ds['test'])}")
+    return ds
+
+# ---------------------------------------------------------------------------
+# Training (GPU)
+# ---------------------------------------------------------------------------
+@app.function(
+    gpu="A10G",  # or "T4", "A100", "H100"
+    volumes={"/data": vol, "/model": MODEL_VOL},
+    secrets=[hf_secret],
+    timeout=3600*4,  # 4 hours
+)
+def train_model():
+    import os, json, torch
+    from datasets import load_from_disk
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from peft import LoraConfig, TaskType
+    from trl import SFTTrainer, SFTConfig
+
+    MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
+    HUB_MODEL_ID = os.environ.get("HUB_MODEL_ID", "Karthik8nitt/parametric-floorplan-generator")
+    OUTPUT_DIR = "/model/floorplan-model"
+
+    print("Loading dataset from /data...")
+    dataset = load_from_disk("/data/floorplan_synthetic_dataset")
+    print(f"Loaded: {len(dataset['train'])} train, {len(dataset['validation'])} val")
+
+    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=4096,
+        bf16=True,
+        gradient_checkpointing=True,
+        report_to="none",
+        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=dataset["train"],
+        eval_dataset=dataset["validation"],
+        peft_config=peft_config,
+        processing_class=tokenizer,
+    )
+
+    print("Starting training...")
+    trainer.train()
+    trainer.save_model(os.path.join(OUTPUT_DIR, "final"))
+    trainer.push_to_hub()
+    print(f"Done! Model pushed to https://huggingface.co/{HUB_MODEL_ID}")
+
+# ---------------------------------------------------------------------------
+# Local entrypoint
+# ---------------------------------------------------------------------------
+@app.local_entrypoint()
+def main():
+    print("Step 1: Generate dataset...")
+    generate_dataset_on_modal.remote()
+    print("Step 2: Train model on GPU...")
+    train_model.remote()
+    print("All done!")