Spaces:
Sleeping
Sleeping
| import pandas as pd | |
| import numpy as np | |
| from sklearn.model_selection import train_test_split | |
| from sklearn.preprocessing import LabelEncoder | |
| import xgboost as xgb | |
| import joblib | |
| import os | |
| import warnings | |
| warnings.filterwarnings('ignore') | |
| class ImpactForecaster: | |
| def __init__(self): | |
| self.model = xgb.XGBRegressor( | |
| n_estimators=150, | |
| learning_rate=0.1, | |
| max_depth=6, | |
| random_state=42 | |
| ) | |
| self.encoders = {} | |
| self.features = ['latitude', 'longitude', 'event_type_encoded', | |
| 'event_cause_encoded', 'priority_encoded', | |
| 'start_hour', 'day_of_week', 'is_weekend'] | |
| self.is_trained = False | |
| def train(self, df): | |
| """Trains the XGBoost model to predict impact duration.""" | |
| print("Training ImpactForecaster...") | |
| cat_cols = ['event_type', 'event_cause', 'priority'] | |
| df_encoded = df.copy() | |
| for col in cat_cols: | |
| le = LabelEncoder() | |
| df_encoded[col] = df_encoded[col].astype(str) | |
| le.fit(df_encoded[col].tolist() + ['unknown']) | |
| self.encoders[col] = le | |
| df_encoded[col + '_encoded'] = le.transform(df_encoded[col]) | |
| X = df_encoded[self.features] | |
| y = np.log1p(df_encoded['impact_duration_mins']) | |
| X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) | |
| self.model.fit( | |
| X_train, y_train, | |
| eval_set=[(X_test, y_test)], | |
| verbose=False | |
| ) | |
| self.is_trained = True | |
| # Save model | |
| os.makedirs('models', exist_ok=True) | |
| joblib.dump(self.model, 'models/forecaster_model.pkl') | |
| joblib.dump(self.encoders, 'models/forecaster_encoders.pkl') | |
| print("Model training complete and saved.") | |
| def predict(self, event_data): | |
| """Predicts impact duration and calculates congestion radius.""" | |
| if not self.is_trained: | |
| self._load_model() | |
| df = pd.DataFrame([event_data]) | |
| for col, le in self.encoders.items(): | |
| if col in df.columns: | |
| val = str(df[col].iloc[0]) | |
| if val not in le.classes_: | |
| val = 'unknown' | |
| df[col + '_encoded'] = le.transform([val]) | |
| else: | |
| df[col + '_encoded'] = le.transform(['unknown']) | |
| for feat in self.features: | |
| if feat not in df.columns: | |
| df[feat] = 0 | |
| X = df[self.features] | |
| log_pred = self.model.predict(X)[0] | |
| duration_mins = np.expm1(log_pred) | |
| # Congestion radius scales with duration and priority | |
| base_radius = 500 | |
| priority_val = str(event_data.get('priority', '')).lower() | |
| if priority_val == 'high': | |
| base_radius = 800 | |
| elif priority_val == 'low': | |
| base_radius = 300 | |
| congestion_radius = min(base_radius + (duration_mins * 6), 4000) | |
| return { | |
| 'predicted_duration_mins': float(duration_mins), | |
| 'congestion_radius_meters': float(congestion_radius) | |
| } | |
| def _load_model(self): | |
| try: | |
| self.model = joblib.load('models/forecaster_model.pkl') | |
| self.encoders = joblib.load('models/forecaster_encoders.pkl') | |
| self.is_trained = True | |
| except FileNotFoundError: | |
| raise RuntimeError("Model not trained yet. Call train() first.") | |