Upload handler.py with huggingface_hub

cf196e9 verified about 1 year ago

10 kB

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)

	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)