Update main.py

main.py

@@ -7,25 +7,22 @@ import torch.nn as nn
 from torch.utils.data import Dataset, DataLoader
 from sklearn.preprocessing import LabelEncoder, MinMaxScaler
 from sklearn.metrics import mean_absolute_error
-from fastapi import FastAPI, Query
+from fastapi import FastAPI
 from pydantic import BaseModel
 
-# ----------------------------
-# CONFIGURATION
-# ----------------------------
-CSV_PATH = "observation.csv"        # Upload this in HF Space
-MODEL_PATH = "traffic_transformer_model_2.pth"
+# -------------------- CONFIG --------------------
+CSV_PATH = "observation.csv"  # place your CSV in the same dir
+MODEL_PATH = "best_model.pt"  # pretrained weights file
 INPUT_LEN = 72
 BATCH_SIZE = 128
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-CONGESTION_THRESHOLDS = {'LOW': 200, 'MEDIUM': 300}
+CONGESTION_THRESHOLDS = {
+    'LOW': 200,
+    'MEDIUM': 300
+}
 
-app = FastAPI(title="Traffic Congestion Predictor", version="1.0")
-
-# ----------------------------
-# HELPERS
-# ----------------------------
+# -------------------- HELPER FUNCS --------------------
 def determine_congestion_level(mean_count, thresholds):
     if mean_count <= thresholds['LOW']:
         return "Low"
@@ -34,12 +31,23 @@ def determine_congestion_level(mean_count, thresholds):
     else:
         return "High"
 
-
+# -------------------- MODEL --------------------
 class TransformerForecaster(nn.Module):
     def __init__(self, input_size, hidden_size=256, num_heads=8, dropout=0.1, seq_len=INPUT_LEN):
         super().__init__()
+
+        # ✅ Adjust number of heads dynamically if needed
+        if input_size % num_heads != 0:
+            for h in [1, 2, 4, 8, 16]:
+                if input_size % h == 0:
+                    num_heads = h
+                    break
+            else:
+                num_heads = 1  # fallback to 1 head
+
         self.pos_embedding = nn.Parameter(torch.randn(1, seq_len, input_size) * 0.01)
         self.layer_norm = nn.LayerNorm(input_size)
+
         encoder_layer = nn.TransformerEncoderLayer(
             d_model=input_size,
             nhead=num_heads,
@@ -63,34 +71,34 @@ class TransformerForecaster(nn.Module):
         x = x[:, -1, :]
         return self.fc(x).squeeze(-1)
 
-
 class TrafficDataset(Dataset):
     def __init__(self, data, input_len=INPUT_LEN, feature_dim=None):
         self.X, self.y = [], []
         for i in range(len(data) - input_len):
-            self.X.append(data[i:i+input_len, :feature_dim])
-            self.y.append(data[i+input_len, -1])
+            self.X.append(data[i:i + input_len, :feature_dim])
+            self.y.append(data[i + input_len, -1])
         self.X = torch.tensor(np.array(self.X), dtype=torch.float32)
         self.y = torch.tensor(np.array(self.y), dtype=torch.float32)
-
     def __len__(self):
         return len(self.X)
-
     def __getitem__(self, idx):
         return self.X[idx], self.y[idx]
 
+# -------------------- FASTAPI --------------------
+app = FastAPI(title="Traffic Congestion Inference API")
+
+class RequestModel(BaseModel):
+    target_place: str
 
-# ----------------------------
-# LOAD MODEL ON STARTUP
-# ----------------------------
 @app.on_event("startup")
-def load_model():
-    global model_loaded, df, le_to, scaler_X, scaler_y, feature_cols
+def load_data_and_model():
+    global df, scaler_X, scaler_y, feature_cols, model_loaded
 
-    print("Loading model and preparing dataset...")
+    print("Loading CSV...")
     df = pd.read_csv(CSV_PATH)
     df = df.drop(columns=["Datetime"], errors="ignore")
 
+    # Parse hours
     def extract_hour(timeslot):
         if pd.isna(timeslot): return None
         try:
@@ -105,9 +113,9 @@ def load_model():
     df.dropna(subset=['StartHour'], inplace=True)
     df['StartHour'] = df['StartHour'].astype(int)
 
+    # Encode and features
     le_to = LabelEncoder()
     df['To_encoded'] = le_to.fit_transform(df['To'])
-
     df['DayOfYear'] = (df['Month'] - 1) * 30 + df['Day']
     df['Day_sin'] = np.sin(2 * np.pi * df['DayOfYear'] / 365)
     df['Day_cos'] = np.cos(2 * np.pi * df['DayOfYear'] / 365)
@@ -122,9 +130,8 @@ def load_model():
     df['MA_3'] = df.groupby('To')['CarCount'].transform(lambda x: x.rolling(3).mean())
     df['EMA_5'] = df.groupby('To')['CarCount'].transform(lambda x: x.ewm(span=5, adjust=False).mean())
     df['ROLL12_mean'] = df.groupby('To')['CarCount'].transform(lambda x: x.rolling(12).mean())
-    df['ROLL12_std']  = df.groupby('To')['CarCount'].transform(lambda x: x.rolling(12).std())
+    df['ROLL12_std'] = df.groupby('To')['CarCount'].transform(lambda x: x.rolling(12).std())
     df['Diff_1'] = df.groupby('To')['CarCount'].diff(1)
-
     df.dropna(inplace=True)
     df.reset_index(drop=True, inplace=True)
 
@@ -139,42 +146,33 @@ def load_model():
     scaler_y = MinMaxScaler()
     df['y_scaled'] = scaler_y.fit_transform(df[['CarCount']])
 
-    # ✅ Recreate the model architecture and load only weights
-    feature_dim = len(feature_cols)
-    model_loaded = TransformerForecaster(input_size=feature_dim)
-    state_dict = torch.load(MODEL_PATH, map_location=DEVICE, weights_only=True)
-    model_loaded.load_state_dict(state_dict)
+    # Load model
+    print("Loading model...")
+    sample_input = len(feature_cols)
+    model_loaded = TransformerForecaster(sample_input)
+    model_loaded.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
     model_loaded.to(DEVICE)
     model_loaded.eval()
 
-    print("✅ Model weights successfully loaded on device:", DEVICE)
-
-# ----------------------------
-# API INPUT / OUTPUT SCHEMA
-# ----------------------------
-class TrafficInput(BaseModel):
-    place: str = Query(..., description="Target location for congestion prediction")
-
-
-# ----------------------------
-# API ROUTES
-# ----------------------------
-@app.get("/")
-def root():
-    return {"message": "Traffic Congestion Predictor API is running 🚦"}
+    print("✅ Model and data loaded successfully.")
 
-@app.post("/predict")
-def predict_congestion(input_data: TrafficInput):
-    place = input_data.place
-
-    df_target = df[df['To'] == place].copy()
+@app.post("/infer")
+def infer(request: RequestModel):
+    target_place = request.target_place
+    df_target = df[df['To'] == target_place].copy()
     if df_target.empty:
-        return {"error": f"No data found for the target location: {place}"}
+        return {"error": f"No data found for target location: {target_place}"}
 
     values = df_target[feature_cols + ['y_scaled']].values
     feature_dim = len(feature_cols)
 
-    test_dataset = TrafficDataset(values, input_len=INPUT_LEN, feature_dim=feature_dim)
+    train_size = int(len(values) * 0.8)
+    test_data = values[train_size:]
+
+    if len(test_data) < INPUT_LEN + 10:
+        return {"error": f"Not enough data for {target_place}"}
+
+    test_dataset = TrafficDataset(test_data, input_len=INPUT_LEN, feature_dim=feature_dim)
     test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)
 
     preds_scaled_all, actuals_scaled_all = [], []
@@ -200,8 +198,9 @@ def predict_congestion(input_data: TrafficInput):
     congestion_level = determine_congestion_level(mean_predicted_car_count, CONGESTION_THRESHOLDS)
 
     return {
-        "target_place": place,
-        "mean_absolute_error": round(mae_200, 3),
-        "mean_predicted_car_count": round(mean_predicted_car_count, 2),
-        "inferred_congestion_level": congestion_level
-    }
+    "target_place": target_place,
+    "mean_absolute_error": float(mae_200),
+    "mean_predicted_car_count": float(mean_predicted_car_count),
+    "congestion_level": str(congestion_level)
+}
+