Upload vaccine_stockout_predictor.py with huggingface_hub

vaccine_stockout_predictor.py

@@ -0,0 +1,218 @@
+# Vaccine Stock-Out Prediction Pipeline
+
+import pandas as pd
+import joblib
+import os
+from typing import Dict, Tuple, Optional
+
+class VaccineStockoutPredictor:
+    """
+    A comprehensive vaccine stock-out prediction system for multiple vaccines across different countries.
+    
+    Supports 8 different vaccine types:
+    - BCG (Bacille Calmette-Guérin)
+    - HepB (Hepatitis B)
+    - bOPV (bivalent Oral Polio Vaccine)
+    - Penta (Pentavalent)
+    - PCV (Pneumococcal Conjugate Vaccine)
+    - Rota (Rotavirus)
+    - IPV (Inactivated Polio Vaccine)
+    - TT/Td/DT (Tetanus Toxoid/Tetanus-Diphtheria)
+    """
+    
+    def __init__(self):
+        """Initialize the predictor with all trained models."""
+        self.models = {}
+        self.reference_data = None
+        self._load_models()
+        self._load_reference_data()
+    
+    def _load_models(self):
+        """Load all trained models from the models directory."""
+        model_mapping = {
+            'BCG_model.joblib': 'BCG',
+            'HepB_model.joblib': 'HepB',
+            'bOPV_model.joblib': 'bOPV',
+            'Penta_model.joblib': 'Penta',
+            'PCV_model.joblib': 'PCV',
+            'Rota_model.joblib': 'Rota',
+            'IPV_model.joblib': 'IPV',
+            'TT_Td_DT_model.joblib': 'TT/Td/DT'
+        }
+        
+        for filename, vaccine_name in model_mapping.items():
+            model_path = os.path.join('models', filename)
+            if os.path.exists(model_path):
+                self.models[vaccine_name] = joblib.load(model_path)
+                print(f"Loaded model for {vaccine_name}")
+            else:
+                print(f"Warning: Model file {filename} not found")
+    
+    def _load_reference_data(self):
+        """Load reference data for store information."""
+        try:
+            self.reference_data = pd.read_csv('reference_data.csv')
+            print(f"Loaded reference data with {len(self.reference_data)} stores")
+        except Exception as e:
+            print(f"Error loading reference data: {e}")
+    
+    def get_available_vaccines(self) -> list:
+        """Return list of available vaccine types."""
+        return list(self.models.keys())
+    
+    def get_available_countries(self) -> list:
+        """Return list of available countries."""
+        if self.reference_data is not None:
+            return sorted(self.reference_data['CountryName'].unique().tolist())
+        return []
+    
+    def get_available_stores(self, country: str, sc_level: str = None) -> list:
+        """Return list of available stores for a country."""
+        if self.reference_data is not None:
+            mask = self.reference_data['CountryName'] == country
+            if sc_level:
+                mask &= self.reference_data['SCLevel'] == sc_level
+            return sorted(self.reference_data[mask]['StoreName'].unique().tolist())
+        return []
+    
+    def predict_stockout_risk(self, 
+                             country_name: str, 
+                             sc_level: str, 
+                             store_name: str, 
+                             vaccine_type: str, 
+                             current_stock: int) -> Dict:
+        """
+        Predict stock-out risk for a specific vaccine at a specific store.
+        
+        Args:
+            country_name: Name of the country
+            sc_level: Supply chain level (Central, Subnational, LD)
+            store_name: Name of the store
+            vaccine_type: Type of vaccine (BCG, HepB, bOPV, Penta, PCV, Rota, IPV, TT/Td/DT)
+            current_stock: Current stock level
+            
+        Returns:
+            Dictionary containing prediction results
+        """
+        
+        # Validate inputs
+        if vaccine_type not in self.models:
+            return {
+                'error': f"Vaccine type '{vaccine_type}' not supported. Available types: {list(self.models.keys())}"
+            }
+        
+        if self.reference_data is None:
+            return {'error': 'Reference data not loaded'}
+        
+        # Find store information
+        store_info = self.reference_data[
+            (self.reference_data['CountryName'] == country_name) &
+            (self.reference_data['SCLevel'] == sc_level) &
+            (self.reference_data['StoreName'] == store_name)
+        ]
+        
+        if len(store_info) == 0:
+            return {
+                'error': f"Store '{store_name}' not found in {country_name} at {sc_level} level"
+            }
+        
+        store_info = store_info.iloc[0]
+        
+        # Prepare input features
+        vaccine_min_col = f'{vaccine_type}_Min'
+        vaccine_max_col = f'{vaccine_type}_Max'
+        
+        if vaccine_min_col not in store_info or vaccine_max_col not in store_info:
+            return {
+                'error': f"Min/Max data not available for {vaccine_type} at this store"
+            }
+        
+        # Calculate utilization
+        min_stock = store_info[vaccine_min_col]
+        max_stock = store_info[vaccine_max_col]
+        
+        if max_stock <= min_stock:
+            utilization = 0.5  # Default value if range is invalid
+        else:
+            utilization = current_stock / (max_stock - min_stock + 1)
+        
+        # Create input data
+        input_data = {
+            'CountryName': country_name,
+            'SCLevel': sc_level,
+            'StoreName': store_name,
+            'Population': store_info['Population'],
+            'DistanceToParent': store_info['DistanceToParent'],
+            'Latitude': store_info['Latitude'],
+            'Longitude': store_info['Longitude'],
+            'Average_Utilization': store_info['Average_Utilization'],
+            f'{vaccine_type}_Utilization': utilization
+        }
+        
+        input_df = pd.DataFrame([input_data])
+        
+        # Make prediction
+        model = self.models[vaccine_type]
+        prediction = model.predict(input_df)[0]
+        probability = model.predict_proba(input_df)[0][1] if hasattr(model, 'predict_proba') else None
+        
+        return {
+            'vaccine_type': vaccine_type,
+            'country': country_name,
+            'sc_level': sc_level,
+            'store': store_name,
+            'current_stock': current_stock,
+            'min_stock': min_stock,
+            'max_stock': max_stock,
+            'utilization': utilization,
+            'stockout_risk': int(prediction),
+            'risk_probability': float(probability) if probability is not None else None,
+            'risk_level': 'High' if prediction == 1 else 'Low',
+            'recommendation': self._get_recommendation(prediction, probability, current_stock, min_stock, max_stock)
+        }
+    
+    def _get_recommendation(self, prediction: int, probability: float, 
+                           current_stock: int, min_stock: int, max_stock: int) -> str:
+        """Generate recommendation based on prediction results."""
+        if prediction == 1:  # High risk
+            if current_stock <= min_stock:
+                return "URGENT: Stock level below minimum. Immediate restocking required."
+            else:
+                return "High risk of stock-out. Consider restocking soon."
+        else:  # Low risk
+            if current_stock >= max_stock:
+                return "Stock level above maximum. Consider redistribution."
+            else:
+                return "Stock level adequate. Monitor regularly."
+    
+    def batch_predict(self, predictions_list: list) -> list:
+        """
+        Perform batch predictions for multiple stores/vaccines.
+        
+        Args:
+            predictions_list: List of dictionaries with prediction parameters
+            
+        Returns:
+            List of prediction results
+        """
+        results = []
+        for pred_params in predictions_list:
+            result = self.predict_stockout_risk(**pred_params)
+            results.append(result)
+        return results
+
+
+# Example usage
+if __name__ == "__main__":
+    predictor = VaccineStockoutPredictor()
+    
+    # Example prediction
+    result = predictor.predict_stockout_risk(
+        country_name="Afghanistan",
+        sc_level="Subnational",
+        store_name="Kabul",
+        vaccine_type="BCG",
+        current_stock=50000
+    )
+    
+    print(json.dumps(result, indent=2))