handler.py

import pandas as pd
import numpy as np
from typing import Dict, Any
import os
import pickle
import re
from gensim.models.doc2vec import Doc2Vec
import json

class EndpointHandler:
    """
    Custom handler for Hugging Face Inference Endpoints.
    Handles Doc2Vec + Classifier content classification model.
    """

    def __init__(self, path: str = "."):
        """
        Initialize the handler.

        Args:
            path (str): Path to the model files
        """
        self.path = path

        # Try to load model metadata - first try new format, then fallback to old format
        new_summary_path = os.path.join(path, "best_model_summary.json")
        old_summary_path = os.path.join(path, "model_summary.json")

        best_model_info = None

        try:
            # Try new format first (best_model_summary.json)
            if os.path.exists(new_summary_path):
                with open(new_summary_path, 'r') as f:
                    summary = json.load(f)
                best_model_info = summary['best_model']
                print(f"🔍 Using new format: best_model_summary.json")

            # Fallback to old format (model_summary.json)
            elif os.path.exists(old_summary_path):
                with open(old_summary_path, 'r') as f:
                    summary = json.load(f)
                best_model_info = summary['top_3_models'][0]  # rank 1 is the best
                print(f"🔍 Using old format: model_summary.json")

            # If no summary file, try to find model files directly
            else:
                print(f"🔍 No summary file found, searching for model files directly...")
                model_files = [f for f in os.listdir(path) if f.endswith('.pkl')]
                doc2vec_files = [f for f in os.listdir(path) if f.endswith('.model')]

                if not model_files or not doc2vec_files:
                    raise Exception(f"No model files found in {path}")

                # Use the first found files (assuming single best model)
                classifier_file = model_files[0]
                doc2vec_file = doc2vec_files[0]

                # Create a minimal model info structure
                best_model_info = {
                    'files': {
                        'classifier': classifier_file,
                        'doc2vec': doc2vec_file
                    },
                    'config': {
                        'vector_size': 100  # Default, will be verified from actual model
                    },
                    'accuracy': 0.0,  # Unknown
                    'algorithm_type': 'Unknown'
                }
                print(f"🔍 Using direct file discovery: {classifier_file}, {doc2vec_file}")

            # Load the classifier
            classifier_path = os.path.join(path, best_model_info['files']['classifier'])
            with open(classifier_path, 'rb') as f:
                self.classifier = pickle.load(f)

            # Load the MATCHING Doc2Vec model (same one used to train this classifier)
            doc2vec_path = os.path.join(path, best_model_info['files']['doc2vec'])
            print(f"🔍 Loading Doc2Vec from: {doc2vec_path}")

            # Check if associated files exist
            for suffix in ['.wv.vectors.npy', '.syn1neg.npy']:
                associated_file = doc2vec_path + suffix
                if os.path.exists(associated_file):
                    print(f"✅ Found associated file: {associated_file}")
                else:
                    print(f"⚠️  Missing associated file: {associated_file}")

            try:
                self.doc2vec_model = Doc2Vec.load(doc2vec_path)
            except Exception as e:
                print(f"❌ Failed to load Doc2Vec model: {str(e)}")
                print(f"❌ Current working directory: {os.getcwd()}")
                print(f"❌ Files in path {path}:")
                if os.path.exists(path):
                    for f in os.listdir(path):
                        print(f"  - {f}")
                raise Exception(f"Failed to load Doc2Vec model from {doc2vec_path}: {str(e)}")

            # Verify feature dimensions match
            expected_features = self.doc2vec_model.vector_size
            print(f"✅ Doc2Vec vector size: {expected_features}")

            # Update config with actual vector size if it was unknown
            if best_model_info['config']['vector_size'] != expected_features:
                print(f"🔄 Updating vector size in config from {best_model_info['config']['vector_size']} to {expected_features}")
                best_model_info['config']['vector_size'] = expected_features

            print(f"✅ Model config: {best_model_info['config']}")
            print(f"✅ Vector dimensions verified: {expected_features} features")

            # Additional debugging - test a sample vector
            test_text = "test sample text for debugging"
            test_vector = self.get_document_vector_for_test(test_text)
            print(f"🔍 Test vector shape: {test_vector.shape}")
            print(f"🔍 Classifier type: {type(self.classifier).__name__}")

            # Store model info for reference
            self.model_info = best_model_info

            classifier_name = best_model_info['files']['classifier']
            algorithm_type = best_model_info.get('algorithm_type', 'Unknown')
            accuracy = best_model_info.get('accuracy', 0.0)

            print(f"✅ Loaded best model: {classifier_name} ({algorithm_type})")
            print(f"✅ Model accuracy: {accuracy:.4f}")
            print(f"✅ Doc2Vec config: {best_model_info['config']}")

        except Exception as e:
            print(f"❌ Error details: {str(e)}")
            print(f"❌ Error type: {type(e)}")
            raise Exception(f"❌ Failed to load models: {str(e)}")

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Process inference request.

        Args:
            data (Dict): Request payload

        Expected input formats:
            Single prediction:
            {
                "inputs": {
                    "content": "text content",
                    "meta_description": "meta text"
                }
            }

        Returns:
            Dict: Prediction results
        """
        try:
            # Extract inputs from request
            inputs = data.get("inputs", {})

            if not inputs:
                return {"error": "Missing 'inputs' in request data"}

            # Handle single request
            if isinstance(inputs, dict):
                # Single prediction
                return self._predict_single(inputs)
            else:
                return {"error": "Invalid input format. Expected dict."}

        except Exception as e:
            return {"error": f"Prediction failed: {str(e)}"}

    def _predict_single(self, inputs: Dict[str, str]) -> Dict[str, Any]:
        """Handle single prediction request"""
        content = inputs.get("content", "")
        meta_description = inputs.get("meta_description", "")

        if not content:
            return {"error": "Missing required field: 'content'"}

        # Ensure inputs are strings
        content = str(content)
        meta_description = str(meta_description)

        try:
            # Combine text features (same as training)
            combined_text = content + ' ' + meta_description

            # Preprocess text for Doc2Vec (same as training preprocessing)
            preprocessed_words = self._preprocess_text(combined_text)

            if not preprocessed_words:
                # Handle empty text case
                doc_vector = np.zeros(self.doc2vec_model.vector_size)
            else:
                # Get document vector from Doc2Vec model
                doc_vector = self.doc2vec_model.infer_vector(preprocessed_words)

            # Reshape for classifier (expects 2D array)
            doc_vector = doc_vector.reshape(1, -1)

            # Get predictions from classifier
            probabilities = self.classifier.predict_proba(doc_vector)[0]

            # Get class names from classifier
            classes = self.classifier.classes_

            # Create results in the same format as original handler
            results = []
            for i, label in enumerate(classes):
                results.append({
                    "label": label,
                    "score": float(probabilities[i])
                })

            return results

        except Exception as e:
            return {"error": f"Single prediction failed: {str(e)}"}

    def _preprocess_text(self, text):
        """Clean and preprocess text for Doc2Vec (same as training)"""
        if pd.isna(text) or not text:
            return []

        # Convert to lowercase and remove special characters
        text = str(text).lower()
        text = re.sub(r'[^a-zA-Z\s]', '', text)
        text = re.sub(r'\s+', ' ', text).strip()

        # Split into words
        words = text.split()

        # Remove very short words
        words = [word for word in words if len(word) > 2]

        return words

    def get_document_vector_for_test(self, text):
        """Get document vector for testing (same as _preprocess_text + infer_vector)"""
        words = self._preprocess_text(text)
        if not words:
            return np.zeros(self.doc2vec_model.vector_size)
        else:
            return self.doc2vec_model.infer_vector(words)


# For local testing
if __name__ == "__main__":
    # Test the handler locally
    handler = EndpointHandler(".")

    # Test single prediction
    test_single = {
        "inputs": {
            "content": "Machine learning models for content attribution analysis using Doc2Vec embeddings",
            "meta_description": "A comprehensive guide to ML-based content classification with neural networks"
        }
    }

    print("=== Single Prediction Test ===")
    result = handler(test_single)
    print(result)