handler.py

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)]