import logging import warnings import os import json import time import torch from typing import Dict, List, Any, Tuple from concurrent.futures import ThreadPoolExecutor, as_completed from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig from accelerate import Accelerator # Configure logging logging.basicConfig( level=logging.DEBUG, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", ) logging.captureWarnings(True) logger = logging.getLogger(__name__) logger.debug("Starting the handler.py script and initializing components.") class InvalidInputError(Exception): """Custom exception for invalid input errors.""" pass class EndpointHandler: """ EndpointHandler class responsible for processing incoming data, generating model outputs, and returning results in a backwards-compatible format with detailed metrics. """ # Path to the fully merged fine-tuned model on Hugging Face path = "imsevolve/llama3-8b-tag-extractor" def __init__(self, model_dir: str = path): """ Initialize the EndpointHandler, setting up the model, tokenizer, and accelerator. """ self.accelerator = Accelerator() hf_token = os.environ.get("HF_TOKEN") if not hf_token: raise RuntimeError("HF_TOKEN environment variable is missing!") quantization_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, ) # Load tokenizer from the merged model self.tokenizer = AutoTokenizer.from_pretrained(model_dir, token=hf_token) self.tokenizer.pad_token = self.tokenizer.eos_token # Load the fully merged model self.model = AutoModelForCausalLM.from_pretrained( model_dir, quantization_config=quantization_config, device_map="auto", token=hf_token ) self.model = self.accelerator.prepare(self.model) logger.info(f"Loaded merged model from {model_dir}") def format_prompt(self, instruction: str) -> str: """ Formats the instruction into the prompt template required by the model. """ prompt_template = """ ### Instruction: {} ### Response: """ return prompt_template.format(instruction) def extract_tags_from_response(self, response: str) -> List[str]: """ Extracts tags from the model's response, cleans them, and ensures they are unique. """ try: # Extract the part after "### Response:" response_part = response.split("### Response:")[-1].strip() # Split into tags, clean, and remove duplicates tags = [] for tag in response_part.split(","): clean_tag = tag.split("/")[0].split("\n")[0].strip() if clean_tag: tags.append(clean_tag) unique_tags = list(set(tags)) return unique_tags except Exception as e: logger.error(f"extract_tags_from_response: Error while parsing response: {e}") return ["extract_tags_from_response: error_extraction"] def tokenize_prompts(self, prompts: List[str]) -> Dict[str, torch.Tensor]: """ Tokenizes a list of prompts for input to the model. """ try: tokenized_batch = self.tokenizer( prompts, return_tensors="pt", padding=True, truncation=True, max_length=512, pad_to_multiple_of=8, add_special_tokens=False, ).to(self.accelerator.device) logger.debug(f"tokenize_prompts: Tokenized batch input_ids shape: {tokenized_batch['input_ids'].shape}") return tokenized_batch except Exception as e: logger.error(f"tokenize_prompts: Error during tokenization: {e}") raise RuntimeError("tokenize_prompts: Tokenization failed") from e def format_and_tokenize_prompts(self, inputs: List[Dict[str, str]]) -> Tuple[Dict[str, Any], List[str]]: """ Formats and tokenizes prompts from the input data. """ prompts = [self.format_prompt(input_dict["prompt"]) for input_dict in inputs] tokenized_batch = self.tokenize_prompts(prompts) dataPointIds = [input_dict["dataPointId"] for input_dict in inputs] return tokenized_batch, dataPointIds def generate_model_outputs(self, tokenized_batch: Dict[str, Any]) -> List[torch.Tensor]: """ Generates model outputs from the tokenized batch. """ try: outputs = self.model.generate( **tokenized_batch, max_new_tokens=54, do_sample=True, temperature=0.7, top_p=0.9, repetition_penalty=1.1, ) return outputs except Exception as e: logger.error(f"generate_model_outputs: Error during model generation: {e}") return [] def process_model_outputs(self, outputs: List[torch.Tensor], dataPointIds: List[str]) -> List[Dict[str, Any]]: """ Processes model outputs, extracts tags, and associates with data point IDs. """ results = [] for idx, output in enumerate(outputs): try: generated_text = self.tokenizer.decode(output, skip_special_tokens=True) tags = self.extract_tags_from_response(generated_text) results.append({ "dataPointId": dataPointIds[idx], "generated_text": tags }) except Exception as e: logger.error(f"Error processing output for dataPointId {dataPointIds[idx]}: {e}") results.append({ "dataPointId": dataPointIds[idx], "error": "Failed to process output" }) return results def process_batch_inputs(self, inputs: List[Dict[str, str]]) -> List[Dict[str, Any]]: """ Processes a batch of inputs: formats, tokenizes, generates outputs, and processes results. """ tokenized_batch, dataPointIds = self.format_and_tokenize_prompts(inputs) outputs = self.generate_model_outputs(tokenized_batch) results = self.process_model_outputs(outputs, dataPointIds) return results def extract_and_validate_inputs(self, data: Dict[str, Any], metrics: Dict[str, Any]) -> List[Dict[str, str]]: """ Extracts and validates inputs from the received data. """ inputs = data.get("inputs", []) metrics["total_inputs"] = len(inputs) if not isinstance(inputs, list) or not inputs: raise InvalidInputError("No inputs provided or inputs are not a list.") return inputs def handle_batching(self, inputs: List[Dict[str, str]], metrics: Dict[str, Any]) -> List[List[Dict[str, str]]]: """ Handles batching of inputs based on the batch size environment variable. """ batch_size = min(int(os.environ.get("BATCH_SIZE", 10)), 50) prompt_batches = self.create_batches(inputs, batch_size) metrics["parallel_batches"] = len(prompt_batches) return prompt_batches def process_batches_in_parallel(self, prompt_batches: List[List[Dict[str, str]]], metrics: Dict[str, Any]) -> List[Dict[str, Any]]: """ Processes batches of inputs in parallel using ThreadPoolExecutor. """ results = [] try: max_workers = min(int(os.environ.get("MAX_WORKERS", 4)), 4) with ThreadPoolExecutor(max_workers=max_workers) as executor: futures = {executor.submit(self.process_batch_inputs, batch): batch for batch in prompt_batches} for future in as_completed(futures): try: torch.cuda.empty_cache() batch_results = future.result() metrics["successfully_processed"] += sum(1 for result in batch_results if "error" not in result) metrics["failed_processed"] += sum(1 for result in batch_results if "error" in result) results.extend(batch_results) except Exception as e: logger.error(f"process_batches_in_parallel: Exception in future result: {e}") results.append({"error": str(e)}) metrics["failed_processed"] += 1 return results except Exception as e: logger.error(f"process_batches_in_parallel: Failed to process batches in parallel: {e}") raise def gather_and_finalize_results(self, results: List[Dict[str, Any]], metrics: Dict[str, Any], start_time: float) -> Dict[str, Any]: """ Gathers results and finalizes metrics. """ metrics["total_processing_time"] = time.time() - start_time return {"results": results, "metrics": metrics} def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: """ Main entry point for processing incoming data. """ metrics = self.initialize_metrics() start_time = time.time() try: inputs = self.extract_and_validate_inputs(data, metrics) prompt_batches = self.handle_batching(inputs, metrics) self.accelerator.wait_for_everyone() results = self.process_batches_in_parallel(prompt_batches, metrics) return self.gather_and_finalize_results(results, metrics, start_time) except InvalidInputError as e: logger.error(f"__call__: Invalid input error: {e}") return {"results": [{"error": str(e)}], "metrics": metrics} except Exception as e: logger.error(f"__call__: Unexpected error: {e}") return {"results": [{"error": str(e)}], "metrics": metrics} @staticmethod def initialize_metrics() -> Dict[str, Any]: """ Initializes the metrics dictionary. """ return { "total_processing_time": 0, "memory_allocated": torch.cuda.memory_allocated(), "memory_reserved": torch.cuda.memory_reserved(), "parallel_batches": 0, "total_inputs": 0, "successfully_processed": 0, "failed_processed": 0 } @staticmethod def create_batches(inputs: List[Dict[str, str]], batch_size: int) -> List[List[Dict[str, str]]]: """ Splits inputs into batches. """ return [inputs[i:i + batch_size] for i in range(0, len(inputs), batch_size)]