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 # Load model metadata to find the best model summary_path = os.path.join(path, "model_summary.json") try: # Load model summary to get best model info with open(summary_path, 'r') as f: summary = json.load(f) # Get the best model (rank 1) best_model_info = summary['top_3_models'][0] # rank 1 is the best # 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}") print(f"✅ Model config: {best_model_info['config']}") # Double-check that vector_size matches config config_vector_size = best_model_info['config']['vector_size'] if expected_features != config_vector_size: raise Exception(f"❌ Vector size mismatch! Doc2Vec model has {expected_features} but config says {config_vector_size}") 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 print(f"✅ Loaded best model: {best_model_info['classifier']} ({best_model_info['algorithm_type']})") print(f"✅ Model accuracy: {best_model_info['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)