import os import torch import logging import time import traceback import json import re from typing import Dict, List, Any, Union, Generator from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer from threading import Thread # Set up logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) logger = logging.getLogger(__name__) class EndpointHandler: def __init__(self, path=""): """ Initialize the model and tokenizer for Phi-4 inference. Args: path (str): Path to the model directory """ # Set default parameters for inference self.max_new_tokens = 1024 # Keep at 1024 to avoid timeouts self.temperature = 0.7 self.top_p = 0.9 self.do_sample = True # Determine if CUDA is available self.device = "cuda" if torch.cuda.is_available() else "cpu" self.dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32 logger.info(f"Initializing model from {path} on {self.device}") try: # Load tokenizer - use original model ID as fallback # This helps with common tokenizer mismatch issues try: self.tokenizer = AutoTokenizer.from_pretrained(path) logger.info(f"Loaded tokenizer from local path") except Exception as e: logger.warning(f"Failed to load tokenizer from local path: {e}") self.tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-4-mini-instruct") logger.info("Loaded tokenizer from microsoft/Phi-4-mini-instruct") # Ensure tokenizer has EOS token set if self.tokenizer.eos_token_id is None: logger.warning("EOS token not set in tokenizer, using default") self.tokenizer.eos_token_id = 199999 # Phi-4's default EOS token # Load model with appropriate settings self.model = AutoModelForCausalLM.from_pretrained( path, torch_dtype=self.dtype, device_map="auto" if self.device == "cuda" else None, trust_remote_code=True ) # Move model to device if CPU if self.device == "cpu": self.model = self.model.to(self.device) # Set model to evaluation mode self.model.eval() # Print diagnostic information logger.info(f"Model loaded on {self.device} using {self.dtype}") logger.info(f"Tokenizer vocabulary size: {len(self.tokenizer)}") logger.info(f"Model vocabulary size: {self.model.config.vocab_size}") logger.info(f"Model embedding size: {self.model.get_input_embeddings().weight.shape}") if len(self.tokenizer) != self.model.config.vocab_size: logger.warning(f"Tokenizer vocab size ({len(self.tokenizer)}) doesn't match model vocab size ({self.model.config.vocab_size})") except Exception as e: logger.error(f"Error during model initialization: {str(e)}") logger.error(traceback.format_exc()) raise def format_prompt_with_system(self, user_message, system_message=None): """ Format the prompt with system and user messages according to Phi-4 format. Args: user_message (str): The user's message system_message (str, optional): The system message/instruction Returns: str: Formatted prompt ready for the model """ # Format using Phi-4's expected chat template: # <|system|> # {system_message} # <|user|> # {user_message} # <|assistant|> if system_message: prompt = f"<|system|>\n{system_message}\n<|user|>\n{user_message}\n<|assistant|>" else: # If no system message, just use user message with assistant tag prompt = f"<|user|>\n{user_message}\n<|assistant|>" logger.info(f"Formatted prompt with {'system message and ' if system_message else ''}user message") return prompt def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: """ Process the input data and generate a response using the Phi-4 model. Args: data (Dict[str, Any]): Input data containing the prompt and generation parameters Returns: Dict[str, Any]: Model response """ start_time = time.time() logger.info(f"Starting request processing") try: # Extract input parameters with defaults if "inputs" not in data: logger.warning("No 'inputs' field in request data") error_msg = "Missing 'inputs' field in request" return self._format_error_response(error_msg) # Track user and system messages user_message = "" system_message = None # Handle different input formats # 1. Direct string input if isinstance(data["inputs"], str): user_message = data["inputs"] system_message = data.get("parameters", {}).get("system_message", None) # 2. Dict with messages format elif isinstance(data["inputs"], dict) and "messages" in data["inputs"]: messages = data["inputs"]["messages"] # Extract system and user messages for prompt formatting for msg in messages: if msg.get("role") == "system": system_message = msg.get("content", "") elif msg.get("role") == "user": user_message = msg.get("content", "") # 3. Direct messages list format elif isinstance(data["inputs"], list): messages = data["inputs"] # Extract system and user messages for prompt formatting for msg in messages: if msg.get("role") == "system": system_message = msg.get("content", "") elif msg.get("role") == "user": user_message = msg.get("content", "") else: logger.warning(f"Unsupported input format: {type(data['inputs'])}") error_msg = "Unsupported input format. Expected string or messages object." return self._format_error_response(error_msg) logger.info(f"Extracted user message length: {len(user_message)} characters") if system_message: logger.info(f"Extracted system message length: {len(system_message)} characters") # Format the prompt with system and user messages prompt = self.format_prompt_with_system(user_message, system_message) parameters = data.get("parameters", {}) logger.info(f"Processing input with {len(prompt)} characters") # Get generation parameters with fallbacks to defaults max_new_tokens = min(parameters.get("max_new_tokens", self.max_new_tokens), 1024) temperature = parameters.get("temperature", self.temperature) top_p = parameters.get("top_p", self.top_p) do_sample = parameters.get("do_sample", self.do_sample) logger.info(f"Generation parameters: max_new_tokens={max_new_tokens}, temperature={temperature}, top_p={top_p}, do_sample={do_sample}") # Manually implement generation to avoid token index errors try: input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.device) logger.info(f"Input tokens shape: {input_ids.shape}") # Create attention mask attention_mask = torch.ones_like(input_ids) # Perform safe generation with error handling for out-of-vocabulary issues response_text = self._safe_generate( input_ids, attention_mask, max_new_tokens, temperature, top_p, do_sample, prompt ) logger.info(f"Response generation completed, text length: {len(response_text) if isinstance(response_text, str) else 'N/A'}") # Format and return response in OpenAI format if isinstance(response_text, str): response_tokens = len(self.tokenizer.encode(response_text)) if response_text else 0 logger.info(f"Response token count: {response_tokens}") return self._format_openai_response( response_text, input_ids.shape[1], response_tokens ) else: return self._format_error_response(f"Error during generation: {response_text}") except RuntimeError as e: logger.error(f"Runtime Error during generation: {str(e)}") logger.error(traceback.format_exc()) return self._format_error_response(f"Error during generation: {str(e)}") except Exception as e: logger.error(f"Unexpected error during request processing: {str(e)}") logger.error(traceback.format_exc()) return self._format_error_response(f"Unexpected error: {str(e)}") finally: duration = time.time() - start_time logger.info(f"Request processing completed in {duration:.2f} seconds") def _complete_sentence(self, text): """Ensure the text ends with a complete sentence""" # If text is already a complete sentence, return it if text.strip().endswith(('.', '!', '?')): return text # Find the last complete sentence end sentences = re.split(r'([.!?])\s+', text) if len(sentences) <= 1: # No complete sentences found, return as is with ellipsis return text + "..." # Reconstruct text up to the last complete sentence result = "" for i in range(len(sentences) - 1): if i % 2 == 0: # Content before punctuation result += sentences[i] else: # Punctuation result += sentences[i] + " " return result.strip() def _safe_generate(self, input_ids, attention_mask, max_new_tokens, temperature, top_p, do_sample, prompt): """Safely generate text handling potential token index errors""" try: with torch.no_grad(): logger.info("Starting safe generation") # Get the input text to exclude from final output input_text = prompt logger.info(f"Input prompt length: {len(input_text)} characters") # Generate one token at a time to avoid index errors # Use a lower absolute maximum to ensure completion max_steps = min(max_new_tokens, 450) # Adjusted down from 500 current_ids = input_ids.clone() logger.info(f"Generating up to {max_steps} tokens") # Keep track of last 5 tokens to detect repetition last_tokens = [] repetition_detected = False for i in range(max_steps): if i % 50 == 0: logger.info(f"Generated {i} tokens so far") # Early termination if we're getting close to the limit to allow for post-processing if i >= max_steps - 50: # Temporarily decode to check if we have a complete response already temp_text = self.tokenizer.decode(current_ids[0], skip_special_tokens=True) if "<|assistant|>" in temp_text: temp_response = temp_text.split("<|assistant|>")[1].strip() # If we have a reasonably complete response, stop early if len(temp_response) > 100 and temp_response.count('.') >= 3: logger.info(f"Early termination at {i} tokens with complete response detected") break # Get logits for next token outputs = self.model( input_ids=current_ids, attention_mask=attention_mask, return_dict=True ) next_token_logits = outputs.logits[:, -1, :] # Apply temperature and sampling if temperature > 0: next_token_logits = next_token_logits / temperature if do_sample: # Apply top_p sampling sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True) cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) # Remove tokens with cumulative probability above the threshold sorted_indices_to_remove = cumulative_probs > top_p # Shift the indices to the right to keep also the first token above the threshold sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() sorted_indices_to_remove[..., 0] = 0 indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) next_token_logits[indices_to_remove] = -float('Inf') # Sample from the filtered distribution probs = torch.softmax(next_token_logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) else: # Take the token with highest probability next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True) # Add the predicted token to the sequence current_ids = torch.cat([current_ids, next_token], dim=-1) attention_mask = torch.cat([attention_mask, torch.ones_like(next_token)], dim=-1) # Add to last tokens list for repetition detection last_tokens.append(next_token.item()) if len(last_tokens) > 5: last_tokens.pop(0) # Check for repetition (if we have at least 5 tokens) if len(last_tokens) >= 5: # Check if all last 5 tokens are the same if len(set(last_tokens)) == 1: logger.warning(f"Repetition detected after {i+1} tokens, stopping generation") repetition_detected = True break # Check if we've generated an EOS token if next_token[0, 0].item() == self.tokenizer.eos_token_id: logger.info(f"EOS token generated after {i+1} tokens") break # Decode the generated sequence generated_text = self.tokenizer.decode(current_ids[0], skip_special_tokens=True) logger.info(f"Decoded generated text: {len(generated_text)} characters") # Return only the newly generated text (after the assistant tag) split_text = generated_text.split("<|assistant|>") if len(split_text) > 1: assistant_response = split_text[1].strip() logger.info(f"Raw assistant response: {len(assistant_response)} characters") # Process the response to ensure complete sentences response_text = self._complete_sentence(assistant_response) logger.info(f"Processed assistant response: {len(response_text)} characters") else: # Fallback if the expected format is not found logger.warning("Could not find assistant tag in generated text") response_text = generated_text return response_text except Exception as e: logger.error(f"Error in _safe_generate: {str(e)}") logger.error(traceback.format_exc()) return f"Generation error: {str(e)}. Please try a simpler input." def _format_openai_response(self, response_text, prompt_tokens, completion_tokens): """Format the response in OpenAI-style format""" try: # Create a response ID response_id = f"phi4-{int(time.time())}" # Build OpenAI-compatible response openai_response = { "id": response_id, "object": "chat.completion", "created": int(time.time()), "model": "phi-4-mini", "choices": [ { "index": 0, "message": { "role": "assistant", "content": response_text }, "finish_reason": "stop" } ], "usage": { "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": prompt_tokens + completion_tokens } } # For compatibility with Hugging Face UI, include the generated_text field openai_response["generated_text"] = response_text logger.info(f"Formatted OpenAI-style response: {len(json.dumps(openai_response))} bytes") return openai_response except Exception as e: logger.error(f"Error formatting OpenAI response: {str(e)}") # Fall back to simple response return {"generated_text": response_text} def _format_error_response(self, error_message): """Format an error response in OpenAI-style format""" try: error_response = { "id": f"phi4-error-{int(time.time())}", "object": "chat.completion", "created": int(time.time()), "model": "phi-4-mini", "choices": [ { "index": 0, "message": { "role": "assistant", "content": f"Error: {error_message}" }, "finish_reason": "error" } ], "usage": { "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0 }, "error": { "message": error_message, "type": "invalid_request_error", "code": "error" } } # For compatibility with Hugging Face UI, include the generated_text field error_response["generated_text"] = f"Error: {error_message}" logger.info(f"Formatted error response: {len(json.dumps(error_response))} bytes") return error_response except Exception as e: logger.error(f"Error formatting error response: {str(e)}") # Fall back to simple error response return {"generated_text": f"Error: {error_message}"} # For local testing if __name__ == "__main__": # Example usage handler = EndpointHandler() # Test with messages format test_with_messages = { "inputs": { "messages": [ {"role": "system", "content": "You are an AI assistant that provides helpful, accurate, and concise information about AI models."}, {"role": "user", "content": "What are the major features of Phi-4?"} ] } } # Run the test result = handler(test_with_messages) print(json.dumps(result, indent=2))