Spaces:
Sleeping
Sleeping
File size: 1,074 Bytes
6b370f6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | # train_model.py
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestRegressor
import joblib, os
def generate_data(n=3000):
np.random.seed(42)
dist = np.random.exponential(400, n)
speed = np.random.uniform(5, 15, n)
tod = np.random.uniform(0, 86400, n)
tod_sin = np.sin(2*np.pi * tod/86400)
tod_cos = np.cos(2*np.pi * tod/86400)
dow = np.random.randint(0, 7, n)
eta = dist/(speed*0.277) + np.random.randint(10,120,n)
df = pd.DataFrame({
"distance_to_stop_m": dist,
"speed_mps": speed,
"tod_sin": tod_sin,
"tod_cos": tod_cos,
"day_of_week": dow,
"eta": eta
})
return df
def train():
os.makedirs("models", exist_ok=True)
df = generate_data()
X = df[["distance_to_stop_m","speed_mps","tod_sin","tod_cos","day_of_week"]]
y = df["eta"]
model = RandomForestRegressor(n_estimators=40)
model.fit(X, y)
joblib.dump(model, "models/model.joblib")
print("Model saved to models/model.joblib")
if __name__ == "__main__":
train()
|