Upload handler.py with huggingface_hub

handler.py

@@ -0,0 +1,207 @@
+import os
+import tempfile
+import pandas as pd
+import logging
+from model import PoultryTFTModel
+import numpy as np
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+class EndpointHandler:
+    def __init__(self, path="."):
+        """
+        Initializes the model handler for a Hugging Face Inference Endpoint.
+        It loads two separate models: one for mortality and one for weight prediction.
+
+        Args:
+            path (str): The path to the model directory where assets are stored.
+        """
+        # --- Environment Variables for Model Configuration ---
+        # Get model filenames from environment variables with defaults.
+        mortality_model_name = os.environ.get("MORTALITY_MODEL_NAME", "ml_models/tft_poultry_mortality.ckpt")
+        mortality_data_name = os.environ.get("MORTALITY_DATA_NAME", "ts_dataset/timeseris_dataset_mortality.pt")
+        weight_model_name = os.environ.get("WEIGHT_MODEL_NAME", "ml_models/tft_poultry_weight.ckpt")
+        weight_data_name = os.environ.get("WEIGHT_DATA_NAME", "ts_dataset/timeseris_dataset_weight.pt")
+
+        # --- Load Mortality Model ---
+        logger.info("--- Loading Mortality Model ---")
+        mortality_model_path = os.path.join(path, mortality_model_name)
+        mortality_data_path = os.path.join(path, mortality_data_name)
+        self.mortality_model = PoultryTFTModel(
+            model_path=mortality_model_path,
+            training_data_path=mortality_data_path,
+            target_column="cummlative_daily_mortality_percentage_ma3"
+        )
+
+        # --- Load Weight Model ---
+        logger.info("\n--- Loading Weight Model ---")
+        weight_model_path = os.path.join(path, weight_model_name)
+        weight_data_path = os.path.join(path, weight_data_name)
+        self.weight_model = PoultryTFTModel(
+            model_path=weight_model_path,
+            training_data_path=weight_data_path,
+            target_column="weight_combined_3"
+        )
+        
+        logger.info("\nModel handler initialized successfully with both models.")
+
+    def __call__(self, data: dict) -> list:
+        """
+        Processes an inference request. For any given input CSV, it predicts BOTH
+        mortality and weight and returns the combined results.
+
+        Args:
+            data (dict): The input data from the request. This is the raw binary CSV content.
+
+        Returns:
+            list: A list of dictionaries with combined prediction results for both targets.
+        """
+        # --- Extract Input Data ---
+        if isinstance(data, dict):
+            csv_bytes = data.pop("inputs", None)
+        else:
+            csv_bytes = data
+
+        if not csv_bytes:
+            return [{"error": "Missing CSV data in payload."}]
+
+        # If the input is a string (from JSON), encode it
+        if isinstance(csv_bytes, str):
+            csv_bytes = csv_bytes.encode('utf-8')
+
+        tmp_path = None
+        try:
+            # Write the raw bytes to a single temporary file for both models to use
+            with tempfile.NamedTemporaryFile(mode='wb', delete=False, suffix=".csv") as tmp:
+                tmp.write(csv_bytes)
+                tmp_path = tmp.name
+            
+            logger.info(f"CSV data written to temporary file: {tmp_path}")
+
+            # --- Helper function for mapping predictions ---
+            def _map_predictions(test_df, predictions_df, col_name):
+                """Maps predictions to the last N days of each cycle based on age."""
+                unique_cycles = test_df['cycle_id'].unique()
+
+
+                pred_values_df = predictions_df
+                cycle_to_predictions = {}
+                for i, cycle_id in enumerate(unique_cycles):
+                    if i < len(pred_values_df):
+                        predictions = pred_values_df.iloc[i].values
+                        cycle_to_predictions[cycle_id] = predictions
+
+                if not cycle_to_predictions:
+                    return pd.DataFrame(columns=['cycle_id', 'age', col_name])
+
+                prediction_horizon = len(next(iter(cycle_to_predictions.values())))
+                
+                last_n_days_data = []
+                for cycle_id in unique_cycles:
+                    if cycle_id in cycle_to_predictions:
+                        cycle_data = test_df[test_df['cycle_id'] == cycle_id].copy().sort_values('age')
+                        last_n_days_data.append(cycle_data.tail(prediction_horizon))
+                
+                if not last_n_days_data:
+                    return pd.DataFrame(columns=['cycle_id', 'age', col_name])
+
+                last_n_days_df = pd.concat(last_n_days_data, ignore_index=True)
+                
+                predictions_list = []
+                for _, row in last_n_days_df.iterrows():
+                    cycle_id = row['cycle_id']
+                    prediction = np.nan
+                    if cycle_id in cycle_to_predictions:
+                        cycle_data_full = test_df[test_df['cycle_id'] == cycle_id].sort_values('age')
+                        ages_in_cycle = cycle_data_full['age'].values
+                        current_age = row['age']
+                        age_positions = np.where(ages_in_cycle == current_age)[0]
+                        
+                        if len(age_positions) > 0:
+                            position_from_end = len(ages_in_cycle) - age_positions[-1] - 1
+                            horizon = len(cycle_to_predictions[cycle_id])
+                            if position_from_end < horizon:
+                                prediction_idx = (horizon - 1) - position_from_end
+                                prediction = cycle_to_predictions[cycle_id][prediction_idx]
+                    predictions_list.append(prediction)
+                
+                last_n_days_df[col_name] = predictions_list
+                return last_n_days_df[['cycle_id', 'age', col_name]]
+
+            # --- Predict Mortality ---
+            logger.info("--- Predicting mortality ---")
+            mortality_test_df, mortality_preds_df = self.mortality_model.predict(tmp_path)
+            mortality_preds_mapped_df = _map_predictions(mortality_test_df, mortality_preds_df, "mortality_predicted")
+
+            # --- Predict Weight ---
+            logger.info("--- Predicting weight ---")
+            weight_test_df, weight_preds_df = self.weight_model.predict(tmp_path)
+            weight_preds_mapped_df = _map_predictions(mortality_test_df, weight_preds_df, "weight_predicted")
+
+            # --- Merge Predictions with Original Data ---
+            logger.info("--- Merging predictions with original data ---")
+            
+            combined_df = mortality_test_df
+
+            predictions_df = pd.merge(
+                mortality_preds_mapped_df,
+                weight_preds_mapped_df,
+                on=['cycle_id', 'age'],
+                how='outer'
+            )
+
+            combined_df = pd.merge(
+                combined_df,
+                predictions_df,
+                on=['cycle_id', 'age'],
+                how='left'
+            )
+            
+            combined_df['mortality_predicted'] = combined_df['mortality_predicted'].where(pd.notna(combined_df['mortality_predicted']), None)
+            combined_df['weight_predicted'] = combined_df['weight_predicted'].where(pd.notna(combined_df['weight_predicted']), None)
+            combined_df["age"] = combined_df["age"].astype(int)
+            # --- Format Output ---
+            logger.info("--- Formatting output ---")
+            output_json = {}
+            for org_code, org_df in combined_df.groupby('organization_code'):
+                output_json[org_code] = {}
+                for farm, farm_df in org_df.groupby('farms'):
+                    house_data = []
+                    # Ensure farm_df is not empty before getting max age
+                    if not farm_df.empty:
+                        max_age = farm_df['age'].max()
+                        for age in range(max_age + 1):
+                            age_data = farm_df[farm_df['age'] == age]
+                            if not age_data.empty:
+                                row = age_data.iloc[0]
+                                house_data.append({
+                                    "age": int(age),
+                                    "weight_actual": row.get('weight_combined_3'),
+                                    "weight_predicted": row.get('weight_predicted'),
+                                    "mortality_actual": row.get('cummlative_daily_mortality_percentage_ma3'),
+                                    "mortality_predicted": row.get('mortality_predicted'),
+                                })
+                            else:
+                                house_data.append({
+                                    "age": int(age),
+                                    "weight_actual": None,
+                                    "weight_predicted": None,
+                                    "mortality_actual": None,
+                                    "mortality_predicted": None,
+                                })
+                    output_json[org_code][farm] = house_data
+
+            result = output_json
+            logger.info("Prediction successful, returning formatted JSON output.")
+
+        except Exception as e:
+            logger.error(f"An error occurred during prediction: {e}", exc_info=True)
+            return [{"error": str(e)}]
+        finally:
+            # Clean up the temporary file
+            if tmp_path and os.path.exists(tmp_path):
+                os.remove(tmp_path)
+                logger.info(f"Temporary file {tmp_path} deleted.")
+        
+        return result