Spaces:
Sleeping
Sleeping
| # 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() | |