| """ |
| 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 |
|
|
| |
| |
| |
| 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) |
|
|
| |
| |
| |
| vol = modal.Volume.from_name("floorplan-data", create_if_missing=True) |
| MODEL_VOL = modal.Volume.from_name("floorplan-model", create_if_missing=True) |
|
|
| |
| |
| |
| hf_secret = modal.Secret.from_name("huggingface-token") |
|
|
| |
| |
| |
| @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 |
|
|
| |
| |
| |
| @app.function( |
| gpu="A10G", |
| volumes={"/data": vol, "/model": MODEL_VOL}, |
| secrets=[hf_secret], |
| timeout=3600*4, |
| ) |
| 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}") |
|
|
| |
| |
| |
| @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!") |
|
|
