Update app.py

app.py

@@ -2,7 +2,7 @@ from fastapi import FastAPI, File, UploadFile, Form, HTTPException
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from contextlib import asynccontextmanager
-import uvicorn
+
 import pandas as pd
 import numpy as np
 import torch
@@ -10,51 +10,88 @@ import json
 import io
 import joblib
 import os
+import sys
+import logging
 from model import DroughtNetLSTM
 from utils import normalize, date_encode, interpolate_nans
 from datetime import datetime
 from typing import List, Optional
 
-# Lifespan event handler (thay thế @app.on_event)
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler(sys.stdout)]
+)
+logger = logging.getLogger(__name__)
+
+# Lifespan event handler
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     global model, scaler_dict, scaler_dict_static, device
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    scaler_dict = joblib.load(os.path.join(os.path.dirname(__file__), "scaler_dict.joblib"))
-    scaler_dict_static = joblib.load(os.path.join(os.path.dirname(__file__), "scaler_dict_static.joblib"))
-    print("Khởi tạo dữ scaler hoàn tất")
-
-    # Define model params
-    time_dim = 20
-    lstm_dim = 256
-    num_layers = 2
-    dropout = 0.15
-    static_dim = 29
-    staticfc_dim = 16
-    hidden_dim = 256
-    output_size = 6
-
-    model = DroughtNetLSTM(
-        time_dim=time_dim,
-        lstm_dim=lstm_dim,
-        num_layers=num_layers,
-        dropout=dropout,
-        static_dim=static_dim,
-        staticfc_dim=staticfc_dim,
-        hidden_dim=hidden_dim,
-        output_size=output_size
-    )
-    model.load_state_dict(torch.load(
-        os.path.join(os.path.dirname(__file__), "best_macro_f1_model.pt"),
-        map_location=device
-    ))
-    model.to(device)
-    model.eval()
-    print("Khởi tạo dữ model hoàn tất")
-
-    yield  # Cho phép app chạy
+    
+    try:
+        logger.info("Starting application initialization")
+        
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        logger.info(f"Using device: {device}")
+
+        # Load scalers with safety measures for version compatibility
+        try:
+            logger.info("Loading scalers")
+            scaler_dict = joblib.load(os.path.join(os.path.dirname(__file__), "scaler_dict.joblib"))
+            scaler_dict_static = joblib.load(os.path.join(os.path.dirname(__file__), "scaler_dict_static.joblib"))
+            logger.info("Scalers loaded successfully")
+        except Exception as e:
+            logger.error(f"Error loading scalers: {str(e)}")
+            # Provide fallback empty dictionaries if loading fails
+            scaler_dict = {}
+            scaler_dict_static = {}
+            logger.warning("Using empty scalers as fallback")
+
+        # Define model params
+        logger.info("Initializing model")
+        time_dim = 20
+        lstm_dim = 256
+        num_layers = 2
+        dropout = 0.15
+        static_dim = 29
+        staticfc_dim = 16
+        hidden_dim = 256
+        output_size = 6
+
+        model = DroughtNetLSTM(
+            time_dim=time_dim,
+            lstm_dim=lstm_dim,
+            num_layers=num_layers,
+            dropout=dropout,
+            static_dim=static_dim,
+            staticfc_dim=staticfc_dim,
+            hidden_dim=hidden_dim,
+            output_size=output_size
+        )
+        
+        try:
+            model_path = os.path.join(os.path.dirname(__file__), "best_macro_f1_model.pt")
+            logger.info(f"Loading model from {model_path}")
+            model.load_state_dict(torch.load(model_path, map_location=device))
+            model.to(device)
+            model.eval()
+            logger.info("Model loaded and initialized successfully")
+        except Exception as e:
+            logger.error(f"Error loading model: {str(e)}")
+            raise  # Re-raise to prevent app from starting with broken model
+            
+        logger.info("Application initialization completed successfully")
+        
+        yield  # Allow app to run
+        
+        logger.info("Application shutdown initiated")
+    except Exception as e:
+        logger.error(f"Critical error during initialization: {str(e)}")
+        # Still yield to allow proper error handling
+        yield
+        logger.info("Application shutdown after initialization error")
 
 app = FastAPI(
     title="Drought Prediction API",
@@ -76,20 +113,32 @@ app.add_middleware(
 async def root():
     return {"message": "Welcome to Drought Prediction API. Use /predict endpoint to make predictions."}
 
+@app.get("/health")
+async def health():
+    """Simple health check endpoint"""
+    return {"status": "ok", "model_loaded": model is not None}
+
 @app.post("/predict")
 async def predict(
     csv_file: UploadFile = File(...),
     x_static: str = Form(...),
 ):
     try:
+        logger.info("Received prediction request")
+        
         # Parse static input
         x_static_list = json.loads(x_static)
         x_static_array = np.array([x_static_list], dtype=np.float32)
+        logger.info(f"Static data shape: {x_static_array.shape}")
 
         # Load and process CSV
         content = await csv_file.read()
-        df = pd.read_csv(io.StringIO(content.decode('utf-8')))
+        df = pd.read_csv(io.StringIO(content.decode('utf-8')), skiprows=26)
+        
+        logger.info(f"Loaded CSV with shape: {df.shape}")
+        
         df = prepare_time_data(df)
+        logger.info("Time data prepared successfully")
 
         # Feature extraction
         float_cols = [
@@ -100,25 +149,36 @@ async def predict(
         features = float_cols + ['sin_day', 'cos_day']
         x_time_array = df[features].to_numpy(dtype=np.float32)
         x_time_array = np.expand_dims(x_time_array, axis=0)
+        logger.info(f"Time features shape: {x_time_array.shape}")
 
         # Normalize
-        x_static_norm, x_time_norm = normalize(
-            x_static_array,
-            x_time_array,
-            scaler_dict=scaler_dict,
-            scaler_dict_static=scaler_dict_static
-        )
+        try:
+            x_static_norm, x_time_norm = normalize(
+                x_static_array,
+                x_time_array,
+                scaler_dict=scaler_dict,
+                scaler_dict_static=scaler_dict_static
+            )
+            logger.info("Data normalized successfully")
+        except Exception as norm_error:
+            logger.error(f"Normalization error: {str(norm_error)}")
+            # Fall back to using unnormalized data if normalization fails
+            logger.warning("Using unnormalized data as fallback")
+            x_static_norm = x_static_array
+            x_time_norm = x_time_array
 
         # To tensors
         x_time_tensor = torch.tensor(x_time_norm).float().to(device)
         x_static_tensor = torch.tensor(x_static_norm).float().to(device)
 
         # Predict
+        logger.info("Running prediction")
         with torch.no_grad():
             output = model(x_time_tensor, x_static_tensor)
             output = torch.clamp(output, min=0.0, max=5.0)
 
         predictions = output.cpu().numpy().tolist()[0]
+        logger.info(f"Prediction completed: {predictions}")
 
         drought_classes = {
             0: "No Drought (D0)",
@@ -141,37 +201,39 @@ async def predict(
             }
         }
 
+        logger.info("Returning prediction result")
         return JSONResponse(content=result)
 
     except Exception as e:
+        logger.error(f"Prediction error: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")
 
 
 def prepare_time_data(df):
-    if 'YEAR' not in df.columns or 'DOY' not in df.columns:
-        if 'date' in df.columns:
-            df['date'] = pd.to_datetime(df['date'])
-            df['YEAR'] = df['date'].dt.year
-            df['DOY'] = df['date'].dt.dayofyear
-        else:
-            raise ValueError("Input CSV must contain either 'date' column or both 'YEAR' and 'DOY' columns")
-
-    if 'date' not in df.columns:
-        df['date'] = pd.to_datetime(df['YEAR'].astype(str) + df['DOY'].astype(str), format="%Y%j")
-
-    df[['sin_day', 'cos_day']] = df['date'].apply(lambda d: pd.Series(date_encode(d)))
-
-    float_cols = [
-        'PRECTOTCORR', 'PS', 'QV2M', 'T2M', 'T2MDEW', 'T2MWET', 'T2M_MAX', 'T2M_MIN', 'T2M_RANGE',
-        'TS', 'WS10M', 'WS10M_MAX', 'WS10M_MIN', 'WS10M_RANGE',
-        'WS50M', 'WS50M_MAX', 'WS50M_MIN', 'WS50M_RANGE',
-    ]
-    for col in float_cols:
-        if col in df.columns and df[col].isna().any():
-            df[col] = interpolate_nans(df[col].values)
-
-    return df
-
-if __name__ == "__main__":
-    port = int(os.environ.get("PORT", 7860))  # Hugging Face Spaces sử dụng cổng 7860
-    uvicorn.run("app:app", host="0.0.0.0", port=port, reload=True)
+    try:
+        if 'YEAR' not in df.columns or 'DOY' not in df.columns:
+            if 'date' in df.columns:
+                df['date'] = pd.to_datetime(df['date'])
+                df['YEAR'] = df['date'].dt.year
+                df['DOY'] = df['date'].dt.dayofyear
+            else:
+                raise ValueError("Input CSV must contain either 'date' column or both 'YEAR' and 'DOY' columns")
+
+        if 'date' not in df.columns:
+            df['date'] = pd.to_datetime(df['YEAR'].astype(str) + df['DOY'].astype(str), format="%Y%j")
+
+        df[['sin_day', 'cos_day']] = df['date'].apply(lambda d: pd.Series(date_encode(d)))
+
+        float_cols = [
+            'PRECTOTCORR', 'PS', 'QV2M', 'T2M', 'T2MDEW', 'T2MWET', 'T2M_MAX', 'T2M_MIN', 'T2M_RANGE',
+            'TS', 'WS10M', 'WS10M_MAX', 'WS10M_MIN', 'WS10M_RANGE',
+            'WS50M', 'WS50M_MAX', 'WS50M_MIN', 'WS50M_RANGE',
+        ]
+        for col in float_cols:
+            if col in df.columns and df[col].isna().any():
+                df[col] = interpolate_nans(df[col].values)
+
+        return df
+    except Exception as e:
+        logger.error(f"Error preparing time data: {str(e)}")
+        raise