Upload 25 files

__script_info__.json

@@ -0,0 +1,8 @@
+{
+    "FileName": "__script_info__.json",
+    "Framework": "huggingface",
+    "ModelId": "none",
+    "Scope": "inference",
+    "Task": "textgeneration1",
+    "Version": "1.1.2"
+}

all_results.json

@@ -0,0 +1,14 @@
+{
+    "epoch": 100.0,
+    "eval_loss": 4.5546875,
+    "eval_runtime": 2.8087,
+    "eval_samples": 11,
+    "eval_samples_per_second": 3.916,
+    "eval_steps_per_second": 0.712,
+    "perplexity": 95.0770390179677,
+    "train_loss": 0.16138767729202907,
+    "train_runtime": 5518.2957,
+    "train_samples": 44,
+    "train_samples_per_second": 0.797,
+    "train_steps_per_second": 0.054
+}

bloom_176b.py

@@ -0,0 +1,146 @@
+import gc
+import logging
+import os
+import traceback
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Union
+
+import deepspeed
+import torch
+from constants import constants
+from djl_python.inputs import Input
+from djl_python.outputs import Output
+from hp_validation import _update_num_beams
+from hp_validation import _validate_payload
+from input import process_input
+from transformers import AutoConfig
+from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
+from transformers import PreTrainedModel
+from transformers import PreTrainedTokenizer
+
+
+TORCH_DTYPE_FROM_STR_MAPPING = {constants.INT8: torch.int8, constants.FP16: torch.float16}
+
+
+class Service(object):
+    """Define Service class for text generation.
+
+    This class aims to have customized model loading and inference methods.
+    """
+
+    def __init__(self):
+        """Initialize model and tokenizer"""
+        self.model: PreTrainedModel = None
+        self.tokenizer: PreTrainedTokenizer = None
+
+    def load_model(self, properties: dict) -> None:
+        """Load Bloom or BloomZ 176B model, tokenizer from disk to memory as instance attributes.
+
+        Args:
+            properties (dict): a python dictionary of model parameters.
+        """
+        tensor_parallel = properties[constants.TENSOR_PARALLEL_DEGREE]
+        deepspeed.init_distributed(constants.NCCL)
+
+        model_location = properties[constants.MODEL_DIR]
+        if constants.MODEL_ID in properties:
+            model_location = properties[constants.MODEL_ID]
+        logging.info(f"Loading model from disk at '{model_location}'.")
+
+        curr_pid = os.getpid()
+        logging.info(f"Tensor_parallel={tensor_parallel}::curr_pid={curr_pid}::")
+
+        tokenizer = AutoTokenizer.from_pretrained(model_location)
+        config = AutoConfig.from_pretrained(model_location)
+        # Construct model with fake meta tensors, later will be replaced during ds-inference ckpt load
+        with deepspeed.OnDevice(dtype=torch.float16, device=constants.META):
+            model = AutoModelForCausalLM.from_config(config, torch_dtype=torch.float16)
+        model = model.eval()
+        torch.cuda.empty_cache()
+
+        # Deepspeed-Inference Loading
+        repo_root = model_location
+        # tp presharded repos come with their own checkpoints config file
+        checkpoints_json = os.path.join(repo_root, constants.DS_INFERENCE_CONFIG_FILE)
+        model = deepspeed.init_inference(
+            model,
+            tensor_parallel={
+                constants.TP_SIZE: tensor_parallel,
+                constants.ENABLED: True,
+                constants.MPU: None,
+                constants.TP_GROUP: None,
+            },
+            base_dir=repo_root,
+            dtype=TORCH_DTYPE_FROM_STR_MAPPING[properties.get("dtype")],
+            max_tokens=constants.MAX_TOKEN_DS_INIT,
+            checkpoint=checkpoints_json,
+            replace_with_kernel_inject=True,
+        )
+        torch.cuda.empty_cache()
+        deepspeed.runtime.utils.see_memory_usage(constants.POST_DS_INFERENCE_INIT, force=True)
+        model = model.module
+        self.model = model
+        self.tokenizer = tokenizer
+
+    def inference(self, inputs) -> Union[List[Dict[str, List]], List[List[Dict[str, List]]]]:
+        """Conduct inference based on inputs.
+
+        Args:
+            inputs (djl_python.inputs.Input): input containing payload and content type.
+        Returns:
+            results (Union[List[Dict[str, List]], List[List[Dict[str, List]]]]): if the length of input string being
+            one, the return output is a list of dictionary, where the length of dictionary corresponds to the number of
+            return sequences; if the length of string being more than one, (i.e., batch inference), the return output is
+            a list of list which inner list contains the one or multiple dictionaries.
+        """
+        try:
+            input_data, model_kwargs = process_input(inputs=inputs, text_input_for_bloom_model=True)
+        except Exception as e:
+            logging.exception(f"Failed to do inference: {e}; {traceback.format_exc()}")
+            results = Output().error((str(e)))
+            return results
+
+        input_tokens = self.tokenizer.batch_encode_plus(
+            input_data, return_tensors=constants.RETURN_TENSOR_TYPE, padding=True
+        )
+        for t in input_tokens:
+            if torch.is_tensor(input_tokens[t]):
+                input_tokens[t] = input_tokens[t].to(torch.cuda.current_device())
+        outputs = self.model.generate(**input_tokens, **model_kwargs)
+        outputs: List[str] = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+
+        if len(outputs) > 1:
+            results: List[List[Dict[str, List]]] = []
+            num_return_seq = model_kwargs.get(constants.NUM_RETURN_SEQUENCES, 1)
+            for i in range(0, len(outputs), num_return_seq):
+                res_tmp: List[dict] = []
+                for j in range(i, i + num_return_seq):
+                    res_tmp.append({constants.GENERATED_TEXT: outputs[j]})
+                results.append(res_tmp)
+        else:
+            results = [{constants.GENERATED_TEXT: outputs}]
+        return results
+
+
+_service = Service()
+
+
+def handle(inputs: Input) -> Optional[Output]:
+    """Define handler method for Bloom 176B model.
+
+    Args:
+         inputs (djl_python.inputs.Input): input containing payload and content type.
+    Returns:
+        outputs (djl_python.inputs.Output): model prediction output.
+    """
+
+    if inputs.is_empty():
+        # Model server makes an empty call to warmup the model on startup
+        _service.load_model(inputs.get_properties())
+        return None
+
+    results = _service.inference(inputs)
+    return Output().add(results)

config.json

@@ -0,0 +1,25 @@
+{
+  "_name_or_path": "/tmp",
+  "architectures": [
+    "GPTNeoXForCausalLM"
+  ],
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "hidden_act": "gelu",
+  "hidden_size": 2560,
+  "initializer_range": 0.02,
+  "intermediate_size": 10240,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 2048,
+  "model_type": "gpt_neox",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 32,
+  "rotary_emb_base": 10000,
+  "rotary_pct": 1.0,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float16",
+  "transformers_version": "4.28.1",
+  "use_cache": false,
+  "use_parallel_residual": false,
+  "vocab_size": 50277
+}

eval_results.json

@@ -0,0 +1,9 @@
+{
+    "epoch": 100.0,
+    "eval_loss": 4.5546875,
+    "eval_runtime": 2.8087,
+    "eval_samples": 11,
+    "eval_samples_per_second": 3.916,
+    "eval_steps_per_second": 0.712,
+    "perplexity": 95.0770390179677
+}

extra_requirements.txt

@@ -0,0 +1 @@
+huggingface/textgeneration1/flash_attn/v1.0.0/flash_attn-1.0.5-cp39-cp39-linux_x86_64.whl

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "transformers_version": "4.28.1"
+}

hp_validation.py

@@ -0,0 +1,162 @@
+import logging
+from typing import Any
+from typing import Dict
+from typing import Union
+
+
+logging.basicConfig(level=logging.INFO)
+
+
+# Possible model parameters
+TEXT_INPUTS = "text_inputs"
+MAX_LENGTH = "max_length"
+NUM_RETURN_SEQUENCES = "num_return_sequences"
+NUM_BEAMS = "num_beams"
+TOP_P = "top_p"
+EARLY_STOPPING = "early_stopping"
+DO_SAMPLE = "do_sample"
+NO_REPEAT_NGRAM_SIZE = "no_repeat_ngram_size"
+TOP_K = "top_k"
+TEMPERATURE = "temperature"
+SEED = "seed"
+MIN_LENGTH = "min_length"
+MIN_NEW_TOKENS = "min_new_tokens"
+MAX_NEW_TOKENS = "max_new_tokens"
+LENGTH_PENALTY = "length_penalty"
+MAX_TIME = "max_time"
+RETURN_FULL_TEXT = "return_full_text"
+STOPPING_CRITERIA = "stopping_criteria"
+
+ALL_PARAM_NAMES = [
+    TEXT_INPUTS,
+    MAX_LENGTH,
+    NUM_RETURN_SEQUENCES,
+    NUM_BEAMS,
+    TOP_P,
+    EARLY_STOPPING,
+    DO_SAMPLE,
+    NO_REPEAT_NGRAM_SIZE,
+    TOP_K,
+    TEMPERATURE,
+    SEED,
+    MIN_LENGTH,
+    MAX_NEW_TOKENS,
+    MIN_NEW_TOKENS,
+    LENGTH_PENALTY,
+    MAX_TIME,
+    RETURN_FULL_TEXT,
+    STOPPING_CRITERIA,
+]
+
+
+# Model parameter ranges
+LENGTH_MIN = 1
+NUM_RETURN_SEQUENCE_MIN = 1
+NUM_BEAMS_MIN = 1
+TOP_P_MIN = 0
+TOP_P_MAX = 1
+NO_REPEAT_NGRAM_SIZE_MIN = 1
+TOP_K_MIN = 0
+TEMPERATURE_MIN = 0
+NEW_TOKENS_MIN = 0
+
+
+def is_list_of_strings(parameter: Any) -> bool:
+    """Return True if the parameter is a list of strings."""
+    if parameter and isinstance(parameter, list):
+        return all(isinstance(elem, str) for elem in parameter)
+    else:
+        return False
+
+
+def _validate_payload(payload: Dict[str, Any]) -> None:
+    """Validate the parameters in the input loads.
+
+    Checks if max_length, num_return_sequences, num_beams, top_p and temprature are in bounds.
+    Checks if do_sample is boolean.
+    Checks max_length, num_return_sequences, num_beams and seed are integers.
+
+    Args:
+        payload: a decoded input payload (dictionary of input parameter and values)
+
+    Raises: ValueError is any of the check fails.
+    """
+    # For all parameters used in text2text generation task, please see
+    # https://huggingface.co/docs/transformers/main_classes/text_generation#transformers.GenerationConfig
+    for param_name in payload:
+        if param_name not in ALL_PARAM_NAMES:
+            raise ValueError(f"Input payload contains an invalid key '{param_name}'. Valid keys are {ALL_PARAM_NAMES}.")
+
+    if TEXT_INPUTS not in payload:
+        raise ValueError(f"Input payload must contain {TEXT_INPUTS} key.")
+
+    for param_name in [MAX_LENGTH, NUM_RETURN_SEQUENCES, NUM_BEAMS, SEED]:
+        if param_name in payload:
+            if type(payload[param_name]) != int:
+                raise ValueError(f"{param_name} must be an integer, got {payload[param_name]}.")
+
+    if MAX_LENGTH in payload:
+        if payload[MAX_LENGTH] < LENGTH_MIN:
+            raise ValueError(f"{MAX_LENGTH} must be at least {LENGTH_MIN}, got {payload[MAX_LENGTH]}.")
+
+    if MIN_LENGTH in payload:
+        if payload[MIN_LENGTH] < LENGTH_MIN:
+            raise ValueError(f"{MIN_LENGTH} must be at least {LENGTH_MIN}, got {payload[MIN_LENGTH]}.")
+
+    if MAX_NEW_TOKENS in payload:
+        if payload[MAX_NEW_TOKENS] < NEW_TOKENS_MIN:
+            raise ValueError(f"{MAX_NEW_TOKENS} must be at least {NEW_TOKENS_MIN}, got {payload[MAX_NEW_TOKENS]}.")
+
+    if MIN_NEW_TOKENS in payload:
+        if payload[MIN_NEW_TOKENS] < NEW_TOKENS_MIN:
+            raise ValueError(f"{MIN_NEW_TOKENS} must be at least {NEW_TOKENS_MIN}, got {payload[MIN_NEW_TOKENS]}.")
+
+    if NUM_RETURN_SEQUENCES in payload:
+        if payload[NUM_RETURN_SEQUENCES] < NUM_RETURN_SEQUENCE_MIN:
+            raise ValueError(
+                f"{NUM_RETURN_SEQUENCES} must be at least {NUM_RETURN_SEQUENCE_MIN}, "
+                f"got {payload[NUM_RETURN_SEQUENCES]}."
+            )
+
+    if NUM_BEAMS in payload:
+        if payload[NUM_BEAMS] < NUM_BEAMS_MIN:
+            raise ValueError(f"{NUM_BEAMS} must be at least {NUM_BEAMS_MIN}, got {payload[NUM_BEAMS]}.")
+
+    if NUM_RETURN_SEQUENCES in payload and NUM_BEAMS in payload:
+        if payload[NUM_RETURN_SEQUENCES] > payload[NUM_BEAMS]:
+            raise ValueError(
+                f"{NUM_BEAMS} must be at least {NUM_RETURN_SEQUENCES}. Instead got "
+                f"{NUM_BEAMS}={payload[NUM_BEAMS]} and {NUM_RETURN_SEQUENCES}="
+                f"{payload[NUM_RETURN_SEQUENCES]}."
+            )
+
+    if TOP_P in payload:
+        if payload[TOP_P] < TOP_P_MIN or payload[TOP_P] > TOP_P_MAX:
+            raise ValueError(f"{TOP_K} must be in range [{TOP_P_MIN},{TOP_P_MAX}], got " f"{payload[TOP_P]}")
+
+    if TEMPERATURE in payload:
+        if payload[TEMPERATURE] < TEMPERATURE_MIN:
+            raise ValueError(
+                f"{TEMPERATURE} must be a float with value at least {TEMPERATURE_MIN}, got " f"{payload[TEMPERATURE]}."
+            )
+
+    if DO_SAMPLE in payload:
+        if type(payload[DO_SAMPLE]) != bool:
+            raise ValueError(f"{DO_SAMPLE} must be a boolean, got {payload[DO_SAMPLE]}.")
+
+    if STOPPING_CRITERIA in payload and not is_list_of_strings(payload[STOPPING_CRITERIA]):
+        raise ValueError(f"{payload[STOPPING_CRITERIA]} must be a list of strings, got {payload[STOPPING_CRITERIA]}")
+
+
+def _update_num_beams(payload: Dict[str, Union[str, float, int]]) -> Dict[str, Union[str, float, int]]:
+    """Add num_beans to the payload if missing and num_return_sequences is present.
+
+    Args:
+        payload (Dict): dictionary of input text and parameters
+    Returns:
+        payload (Dict): payload with number of beams updated
+    """
+
+    if NUM_RETURN_SEQUENCES in payload and NUM_BEAMS not in payload:
+        payload[NUM_BEAMS] = payload[NUM_RETURN_SEQUENCES]
+    return payload

inference.py

@@ -0,0 +1,62 @@
+import logging
+from typing import Any
+from typing import Dict
+from typing import Optional
+from typing import Union
+
+from djl_python.deepspeed import DeepSpeedService
+from djl_python.inputs import Input
+from djl_python.outputs import Output
+from hp_validation import _update_num_beams
+from hp_validation import _validate_payload
+from input import process_input
+from transformers import set_seed
+
+
+class DeepSpeedServiceTextGeneration(DeepSpeedService):
+    """Define subclass DeepSpeedServiceTextGeneration.
+
+    This class aims to have customized inference function.
+    For the definition of DeepSpeedService, see
+    https://github.com/deepjavalibrary/djl-serving/blob/master/engines/python/setup/djl_python/deepspeed.py
+    """
+
+    def inference(self, inputs: Input) -> Output:
+        """Define customized inference method to have hyperparameter validation for text generation task.
+
+        Args:
+             inputs (djl_python.inputs.Input): input containing payload and content type.
+        Returns:
+            outputs (djl_python.inputs.Output): model prediction output.
+        """
+
+        try:
+            input_data, model_kwargs = process_input(inputs=inputs, text_input_for_bloom_model=False)
+            result = self.pipeline(input_data, **model_kwargs)
+            outputs = Output()
+            outputs.add(result)
+        except Exception as e:
+            logging.exception("Failed to do inference")
+            outputs = Output().error((str(e)))
+        return outputs
+
+
+_service = DeepSpeedServiceTextGeneration()
+
+
+def handle(inputs: Input) -> Optional[Output]:
+    """Define customized inference method to have hyperparameter validation for text generation task
+
+    Args:
+         inputs (djl_python.inputs.Input): input containing payload and content type.
+    Returns:
+        outputs (djl_python.inputs.Output): model prediction output.
+    """
+
+    if not _service.initialized:
+        _service.initialize(inputs.get_properties())
+
+    if inputs.is_empty():
+        return None
+
+    return _service.inference(inputs)

inference_helper.py

@@ -0,0 +1,53 @@
+from typing import Dict
+from typing import List
+from typing import Union
+
+import torch
+from constants import constants
+from djl_python import Input
+from djl_python import Output
+from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
+
+
+def inference_helper_model_tokenizer(
+    input_data: List, model: AutoModelForCausalLM, tokenizer: AutoTokenizer, content_type: str, model_kwargs: Dict
+) -> Union[List[Dict[str, List]], List[List[Dict[str, List]]]]:
+    """Conduct inference based on inputs.
+
+    Args:
+        input_data (list):  a python list of string.
+        model (AutoModelForCausalLM): model for doing the inference
+        tokenizer (AutoTokenizer): tokenizer for the inference
+        content_type (str): request content type
+        model_kwargs (dict): a python dictionary of parameters.
+
+    Returns:
+        results (Union[List[Dict[str, List]], List[List[Dict[str, List]]]]): if the length of input list being
+        one, the return output is a list of dictionary, where the length of dictionary corresponds to the number of
+        return sequences; if the length of input list being more than one, (i.e., batch inference), the return
+        output is a list of list where inner list contains one or multiple dictionaries.
+    """
+
+    input_tokens = tokenizer.batch_encode_plus(input_data, return_tensors=constants.RETURN_TENSOR_TYPE, padding=True)
+    for t in input_tokens:
+        if torch.is_tensor(input_tokens[t]):
+            input_tokens[t] = input_tokens[t].to(torch.cuda.current_device())
+    outputs = model.generate(**input_tokens, **model_kwargs)
+    outputs: List[str] = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+
+    if len(outputs) > 1:
+        results: List[List[Dict[str, List]]] = []
+        num_return_seq = model_kwargs.get(constants.NUM_RETURN_SEQUENCES, 1)
+        for i in range(0, len(outputs), num_return_seq):
+            res_tmp: List[dict] = []
+            for j in range(i, min(i + num_return_seq, len(outputs))):
+                res_tmp.append({constants.GENERATED_TEXT: outputs[j]})
+            results.append(res_tmp)
+    else:
+        if content_type == constants.APPLICATION_X_TEXT:
+            results = [{constants.GENERATED_TEXT: outputs}]
+        else:
+            results = [[{constants.GENERATED_TEXT: outputs}]]
+
+    return results

inference_mpt.py

@@ -0,0 +1,104 @@
+import os
+from pathlib import Path
+from typing import Any
+from typing import Dict
+
+import torch
+from constants import constants
+from djl_python.inputs import Input
+from djl_python.outputs import Output
+from sagemaker_jumpstart_huggingface_script_utilities.djl_python.dtypes import get_torch_dtype_from_str
+from sagemaker_jumpstart_huggingface_script_utilities.djl_python.handle import create_handle
+from sagemaker_jumpstart_huggingface_script_utilities.djl_python.inference.textgeneration import format_djl_output
+from sagemaker_jumpstart_huggingface_script_utilities.djl_python.inference.textgeneration import generate_text
+from sagemaker_jumpstart_huggingface_script_utilities.djl_python.inference.textgeneration import (
+    model_output_to_batch_output,
+)
+from sagemaker_jumpstart_huggingface_script_utilities.djl_python.inference.textgeneration import process_input
+from sagemaker_jumpstart_huggingface_script_utilities.payload.stopping_criteria import (
+    add_stopping_criteria_to_model_kwargs,
+)
+from sagemaker_jumpstart_script_utilities.subprocess import run_with_error_handling
+from transformers import AutoConfig
+from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
+
+
+FLASH_ATTENTION_WHEEL_FILENAME = "flash_attn-1.0.5-cp39-cp39-linux_x86_64.whl"
+MAXIMUM_INPUT_SEQUENCE_LENGTH = "MAXIMUM_INPUT_SEQUENCE_LENGTH"
+DTYPE = "dtype"
+MAX_SEQ_LEN = "max_seq_len"
+ATTN_IMPL = "attn_impl"
+TRITON = "triton"
+LIB = "lib"
+EXTRA_DEPENDENCIES = "extra_dependencies"
+EXTRA_DEPENDENCIES_PATH = Path(__file__).parent / LIB / EXTRA_DEPENDENCIES
+
+# As of the time of this script creation, the maintainers of the flash-attn package do not distribute a bdist. To
+# avoid an approximately 10 minute source build, a pre-built wheel is hosted as an extra dependency.
+run_with_error_handling(["pip", "install", str(EXTRA_DEPENDENCIES_PATH / FLASH_ATTENTION_WHEEL_FILENAME)])
+
+
+class MptPythonServiceTextGeneration:
+    """A service object for the MPT model family using the DJL Python engine.
+
+    This set of MPT models is not compatible with the T4 GPUs on the G4 instance family because some Triton kernels are
+    not compatible with architectures older than Ampere
+    ([reference](https://github.com/microsoft/DeepSpeed-MII/issues/170#issuecomment-1526277566)).
+
+    This set of MPT models uses the [FlashAttention](https://arxiv.org/pdf/2205.14135.pdf) mechanism with implementation
+    that requires the Python environment to include the following libraries:
+    - [flash-attn](https://github.com/HazyResearch/flash-attention)
+    - [triton](https://github.com/openai/triton)
+    - [einops](https://github.com/arogozhnikov/einops)
+    """
+
+    def __init__(self) -> None:
+        """Set initialization flag to False, model to be initialized upon invocation of initialize method."""
+        self.initialized = False
+
+    def initialize(self, properties: Dict[str, Any]) -> None:
+        """Initialize the MPT model and tokenizer.
+
+        This model contains custom Python scripts that are not yet included in the transformers library. Therefore, the
+        non-standard `trust_remote_code=True` must be set when loading config and the model. The referenced remote code
+        exists within the model artifact, which has been manually inspected for malicious code, downloaded with commit
+        revision hash checking, and prepackaged within the script tarball. Therefore, the "remote" code is hosted on S3
+        via the SageMaker JumpStart deployment pipeline. This script will not perform requests to the public internet.
+        Due to these steps and the fact that SageMaker JumpStart deploys models with network isolation enabled by
+        default, these scripts are deemed safe to enable to `trust_remote_code=True` flag.
+
+        Additionally, the accelerate library currently does not support the MPT model type.
+        """
+        model_dir = properties[constants.MODEL_DIR]
+
+        # View docstring for details on this `trust_remote_code=True` flag.
+        config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
+        config.attn_config[ATTN_IMPL] = TRITON
+        maximum_sequence_length = os.environ.get(MAXIMUM_INPUT_SEQUENCE_LENGTH)
+        if maximum_sequence_length is not None:
+            config.update({MAX_SEQ_LEN: int(maximum_sequence_length)})
+
+        torch_dtype = get_torch_dtype_from_str(properties.get(DTYPE))
+
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_dir, config=config, torch_dtype=torch_dtype, trust_remote_code=True
+        )
+        self.model.to(device="cuda:0")
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        self.initialized = True
+
+    @format_djl_output
+    def inference(self, inputs: Input) -> Output:
+        """Perform inference for text generation task."""
+        input_data, model_kwargs = process_input(inputs, input_data_as_list=False)
+        model_kwargs = add_stopping_criteria_to_model_kwargs(model_kwargs, self.tokenizer)
+        model_output = generate_text(self.model, self.tokenizer, input_data, model_kwargs)
+        return model_output_to_batch_output(model_output, model_kwargs)
+
+
+_service = MptPythonServiceTextGeneration()
+handle = create_handle(_service)

input.py

@@ -0,0 +1,76 @@
+import logging
+from typing import Dict
+from typing import List
+from typing import Union
+
+from constants import constants
+from djl_python.inputs import Input
+from hp_validation import _update_num_beams
+from hp_validation import _validate_payload
+from transformers import set_seed
+
+
+def format_input_for_task(input_values: Union[str, List[str]]) -> List[str]:
+    """Format input string into a list for text generation task
+
+    Args:
+        input_values: either a text string or list of text string.
+    Returns:
+        input_values (list): list of text string.
+    """
+    if not isinstance(input_values, list):
+        input_values = [input_values]
+    return input_values
+
+
+def process_input(inputs: Input, text_input_for_bloom_model: bool = True) -> Union[List, Dict]:
+    """Process input based on content type.
+
+    Parse the input based on Content-Type of application/json or application/x-text
+
+    Args:
+        inputs (djl_python.inputs.Input): input containing payload and content type.
+        text_input_for_bloom_model: whether text string need special handling for bloom model
+    Returns:
+        input_data (list):  a python list of string.
+        model_kwargs (dict): a python dictionary of parameters.
+    """
+    content_type = inputs.get_property("Content-Type")
+    model_kwargs = {}
+    if content_type == constants.APPLICATION_JSON:
+        try:
+            json_input = inputs.get_as_json()
+
+            _validate_payload(json_input)
+            json_input = _update_num_beams(json_input)
+            if constants.SEED in json_input:
+                set_seed(json_input[constants.SEED])
+                del json_input[constants.SEED]
+
+            if isinstance(json_input, dict):
+                input_data = format_input_for_task(json_input.pop(constants.TEXT_INPUTS))
+                model_kwargs = json_input
+            else:
+                input_data = json_input
+        except Exception:
+            logging.exception(
+                f"Failed to parse input payload. For content_type={constants.APPLICATION_JSON}, input "
+                f"payload must be a json encoded dictionary with keys {constants.ALL_PARAM_NAMES}."
+            )
+            raise
+    elif content_type == constants.APPLICATION_X_TEXT:
+        try:
+            if text_input_for_bloom_model:
+                input_data = format_input_for_task(inputs.get_as_string())
+            else:
+                input_data = inputs.get_as_string()
+        except Exception:
+            logging.exception(
+                f"Failed to parse input payload. For content_type={constants.APPLICATION_X_TEXT}, input "
+                f"payload must be a string encoded in utf-8 format."
+            )
+            raise
+    else:
+        raise ValueError('{{"error": "unsupported content type {}"}}'.format(content_type or "unknown"))
+
+    return input_data, model_kwargs

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:63818b91e6ae42613ead2d886f433c31f8996fd5c33c0b37e158a181bab3bb23
+size 5684519513

red_pajama.py

@@ -0,0 +1,131 @@
+import logging
+import traceback
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Union
+
+import torch
+from constants import constants
+from djl_python import Input
+from djl_python import Output
+from inference_helper import inference_helper_model_tokenizer
+from input import process_input
+from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
+from transformers import PreTrainedModel
+from transformers import PreTrainedTokenizer
+from transformers import PreTrainedTokenizerBase
+from transformers import StoppingCriteria
+from transformers import StoppingCriteriaList
+
+
+class StopWordsCriteria(StoppingCriteria):
+    """A text generation stopping criteria when output sequence contains any specified stop words."""
+
+    def __init__(self, tokenizer: PreTrainedTokenizerBase, stop_words: List[str]) -> None:
+        """Initialize stopping criteria."""
+        self._tokenizer = tokenizer
+        self._stop_words = stop_words
+        self._partial_result = ""
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> bool:
+        """Return True if any stop word is present in generated sequence."""
+        text = self._tokenizer.decode(input_ids[0, -1])
+        self._partial_result += text
+        for stop_word in self._stop_words:
+            if stop_word in self._partial_result:
+                return True
+        return False
+
+
+def create_stopping_criteria_list(tokenizer: PreTrainedTokenizerBase, stop_words: List[str]) -> StoppingCriteriaList:
+    """Create a StoppingCriteriaList to be used as generation payload input."""
+    stop_criteria = StopWordsCriteria(tokenizer, stop_words)
+    return StoppingCriteriaList([stop_criteria])
+
+
+class Service(object):
+    """Define Service class for text generation.
+
+    This class aims to have customized model loading and inference methods.
+    """
+
+    def __init__(self):
+        """Initialize model and tokenizer"""
+        self.model: Optional[PreTrainedModel] = None
+        self.tokenizer: Optional[PreTrainedTokenizer] = None
+
+    def load_model(self, properties: dict) -> None:
+        """Load model, tokenizer from disk to memory as instance attributes.
+
+        Args:
+            properties (dict): a python dictionary of model parameters.
+        """
+
+        model_location = properties["model_dir"]
+        logging.info(f"Loading model from {model_location}")
+
+        tokenizer = AutoTokenizer.from_pretrained(model_location)
+        tokenizer.pad_token = tokenizer.eos_token
+        model = AutoModelForCausalLM.from_pretrained(model_location, torch_dtype=torch.float16, device_map="auto")
+        model.requires_grad_(False)
+        model.eval()
+
+        self.model = model
+        self.tokenizer = tokenizer
+
+    def modify_model_kwargs(self, model_kwargs: Dict[str, Any]) -> None:
+        """Provide script-based modification of model kwargs.
+
+        By default, this method injects stopping criteria into text generation. Since payload parameters are serialized
+        when invoking the endpoint, this method initializes any non-serializable kwargs for text generation.
+        """
+        if constants.STOPPING_CRITERIA in model_kwargs:
+            stopping_criteria = model_kwargs.pop(constants.STOPPING_CRITERIA)
+            stopping_criteria_list = create_stopping_criteria_list(self.tokenizer, stopping_criteria)
+            model_kwargs[constants.STOPPING_CRITERIA] = stopping_criteria_list
+
+    def inference(self, inputs: Input) -> Union[List[Dict[str, List]], List[List[Dict[str, List]]]]:
+        """Conduct inference based on inputs.
+
+        Args:
+            inputs (djl_python.inputs.Input): input containing payload and content type.
+        Returns:
+            results (Union[List[Dict[str, List]], List[List[Dict[str, List]]]]): if the length of input list being
+            one, the return output is a list of dictionary, where the length of dictionary corresponds to the number of
+            return sequences; if the length of input list being more than one, (i.e., batch inference), the return
+            output is a list of list where inner list contains one or multiple dictionaries.
+        """
+        try:
+            input_data, model_kwargs = process_input(inputs=inputs, text_input_for_bloom_model=True)
+            self.modify_model_kwargs(model_kwargs)
+            content_type = inputs.get_property("Content-Type")
+
+            return inference_helper_model_tokenizer(input_data, self.model, self.tokenizer, content_type, model_kwargs)
+        except Exception as e:
+            logging.exception(f"Failed to do inference: {e}; {traceback.format_exc()}")
+            results = Output().error((str(e)))
+            return results
+
+
+_service = Service()
+
+
+def handle(inputs: Input) -> Optional[Output]:
+    """Define handler method for model.
+
+    Args:
+         inputs (djl_python.inputs.Input): input containing payload and content type.
+    Returns:
+        outputs (djl_python.inputs.Output): model prediction output.
+    """
+
+    if inputs.is_empty():
+        # Model server makes an empty call to warmup the model on startup
+        _service.load_model(inputs.get_properties())
+        return None
+
+    results = _service.inference(inputs)
+    return Output().add(results)

requirements.txt

@@ -0,0 +1,5 @@
+/opt/ml/model/lib/einops/einops-0.6.1-py3-none-any.whl
+/opt/ml/model/lib/ninja/ninja-1.11.1-py2.py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.whl
+/opt/ml/model/lib/triton/triton-2.0.0.dev20221202-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+/opt/ml/model/lib/sagemaker_jumpstart_huggingface_script_utilities/sagemaker_jumpstart_huggingface_script_utilities-1.0.4-py2.py3-none-any.whl
+/opt/ml/model/lib/sagemaker_jumpstart_script_utilities/sagemaker_jumpstart_script_utilities-1.1.7-py2.py3-none-any.whl

script_requirements.txt

@@ -0,0 +1,8 @@
+sagemaker_jumpstart_huggingface_script_utilities==1.0.4
+sagemaker_jumpstart_script_utilities==1.1.7
+einops==0.6.1 \
+    --hash=sha256:99149e46cc808956b174932fe563d920db4d6e5dadb8c6ecdaa7483b7ef7cfc3
+triton==2.0.0.dev20221202 \
+    --hash=sha256:42b3c00ffdc6311ee60318df72a5bddccc9e45951a5837c19d19ef93c72e8186
+ninja==1.11.1 \
+    --hash=sha256:817e2aee2a4d28a708a67bcfba1817ae502c32c6d8ef80e50d63b0f23adf3a08

serving.properties

@@ -0,0 +1,2 @@
+engine=Python
+option.entryPoint=red_pajama.py

special_tokens_map.json

@@ -0,0 +1,5 @@
+{
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>"
+}

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "add_prefix_space": false,
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 2048,
+  "tokenizer_class": "GPTNeoXTokenizer",
+  "unk_token": "<|endoftext|>"
+}

train_results.json

@@ -0,0 +1,8 @@
+{
+    "epoch": 100.0,
+    "train_loss": 0.16138767729202907,
+    "train_runtime": 5518.2957,
+    "train_samples": 44,
+    "train_samples_per_second": 0.797,
+    "train_steps_per_second": 0.054
+}

trainer_state.json

@@ -0,0 +1,325 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 100.0,
+  "global_step": 300,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 3.33,
+      "learning_rate": 4.061954955173073e-06,
+      "loss": 2.0922,
+      "step": 10
+    },
+    {
+      "epoch": 6.67,
+      "learning_rate": 5.28472523771611e-06,
+      "loss": 1.5474,
+      "step": 20
+    },
+    {
+      "epoch": 6.67,
+      "eval_loss": 2.08203125,
+      "eval_runtime": 2.8721,
+      "eval_samples_per_second": 3.83,
+      "eval_steps_per_second": 0.696,
+      "step": 20
+    },
+    {
+      "epoch": 10.0,
+      "learning_rate": 6e-06,
+      "loss": 0.8877,
+      "step": 30
+    },
+    {
+      "epoch": 13.33,
+      "learning_rate": 6e-06,
+      "loss": 0.2622,
+      "step": 40
+    },
+    {
+      "epoch": 13.33,
+      "eval_loss": 2.9765625,
+      "eval_runtime": 2.8096,
+      "eval_samples_per_second": 3.915,
+      "eval_steps_per_second": 0.712,
+      "step": 40
+    },
+    {
+      "epoch": 16.67,
+      "learning_rate": 6e-06,
+      "loss": 0.0353,
+      "step": 50
+    },
+    {
+      "epoch": 20.0,
+      "learning_rate": 6e-06,
+      "loss": 0.0066,
+      "step": 60
+    },
+    {
+      "epoch": 20.0,
+      "eval_loss": 3.96484375,
+      "eval_runtime": 2.8073,
+      "eval_samples_per_second": 3.918,
+      "eval_steps_per_second": 0.712,
+      "step": 60
+    },
+    {
+      "epoch": 23.33,
+      "learning_rate": 6e-06,
+      "loss": 0.0022,
+      "step": 70
+    },
+    {
+      "epoch": 26.67,
+      "learning_rate": 6e-06,
+      "loss": 0.001,
+      "step": 80
+    },
+    {
+      "epoch": 26.67,
+      "eval_loss": 4.140625,
+      "eval_runtime": 2.8092,
+      "eval_samples_per_second": 3.916,
+      "eval_steps_per_second": 0.712,
+      "step": 80
+    },
+    {
+      "epoch": 30.0,
+      "learning_rate": 6e-06,
+      "loss": 0.0009,
+      "step": 90
+    },
+    {
+      "epoch": 33.33,
+      "learning_rate": 6e-06,
+      "loss": 0.0006,
+      "step": 100
+    },
+    {
+      "epoch": 33.33,
+      "eval_loss": 4.22265625,
+      "eval_runtime": 2.8076,
+      "eval_samples_per_second": 3.918,
+      "eval_steps_per_second": 0.712,
+      "step": 100
+    },
+    {
+      "epoch": 36.67,
+      "learning_rate": 6e-06,
+      "loss": 0.0005,
+      "step": 110
+    },
+    {
+      "epoch": 40.0,
+      "learning_rate": 6e-06,
+      "loss": 0.0004,
+      "step": 120
+    },
+    {
+      "epoch": 40.0,
+      "eval_loss": 4.2890625,
+      "eval_runtime": 2.8093,
+      "eval_samples_per_second": 3.916,
+      "eval_steps_per_second": 0.712,
+      "step": 120
+    },
+    {
+      "epoch": 43.33,
+      "learning_rate": 6e-06,
+      "loss": 0.0003,
+      "step": 130
+    },
+    {
+      "epoch": 46.67,
+      "learning_rate": 6e-06,
+      "loss": 0.0003,
+      "step": 140
+    },
+    {
+      "epoch": 46.67,
+      "eval_loss": 4.3515625,
+      "eval_runtime": 2.8091,
+      "eval_samples_per_second": 3.916,
+      "eval_steps_per_second": 0.712,
+      "step": 140
+    },
+    {
+      "epoch": 50.0,
+      "learning_rate": 6e-06,
+      "loss": 0.0003,
+      "step": 150
+    },
+    {
+      "epoch": 53.33,
+      "learning_rate": 6e-06,
+      "loss": 0.0003,
+      "step": 160
+    },
+    {
+      "epoch": 53.33,
+      "eval_loss": 4.390625,
+      "eval_runtime": 2.8088,
+      "eval_samples_per_second": 3.916,
+      "eval_steps_per_second": 0.712,
+      "step": 160
+    },
+    {
+      "epoch": 56.67,
+      "learning_rate": 6e-06,
+      "loss": 0.0003,
+      "step": 170
+    },
+    {
+      "epoch": 60.0,
+      "learning_rate": 6e-06,
+      "loss": 0.0003,
+      "step": 180
+    },
+    {
+      "epoch": 60.0,
+      "eval_loss": 4.42578125,
+      "eval_runtime": 2.808,
+      "eval_samples_per_second": 3.917,
+      "eval_steps_per_second": 0.712,
+      "step": 180
+    },
+    {
+      "epoch": 63.33,
+      "learning_rate": 6e-06,
+      "loss": 0.0002,
+      "step": 190
+    },
+    {
+      "epoch": 66.67,
+      "learning_rate": 6e-06,
+      "loss": 0.0002,
+      "step": 200
+    },
+    {
+      "epoch": 66.67,
+      "eval_loss": 4.453125,
+      "eval_runtime": 2.8089,
+      "eval_samples_per_second": 3.916,
+      "eval_steps_per_second": 0.712,
+      "step": 200
+    },
+    {
+      "epoch": 70.0,
+      "learning_rate": 6e-06,
+      "loss": 0.0003,
+      "step": 210
+    },
+    {
+      "epoch": 73.33,
+      "learning_rate": 6e-06,
+      "loss": 0.0003,
+      "step": 220
+    },
+    {
+      "epoch": 73.33,
+      "eval_loss": 4.47265625,
+      "eval_runtime": 2.8074,
+      "eval_samples_per_second": 3.918,
+      "eval_steps_per_second": 0.712,
+      "step": 220
+    },
+    {
+      "epoch": 76.67,
+      "learning_rate": 6e-06,
+      "loss": 0.0002,
+      "step": 230
+    },
+    {
+      "epoch": 80.0,
+      "learning_rate": 6e-06,
+      "loss": 0.0002,
+      "step": 240
+    },
+    {
+      "epoch": 80.0,
+      "eval_loss": 4.49609375,
+      "eval_runtime": 2.8069,
+      "eval_samples_per_second": 3.919,
+      "eval_steps_per_second": 0.713,
+      "step": 240
+    },
+    {
+      "epoch": 83.33,
+      "learning_rate": 6e-06,
+      "loss": 0.0002,
+      "step": 250
+    },
+    {
+      "epoch": 86.67,
+      "learning_rate": 6e-06,
+      "loss": 0.0002,
+      "step": 260
+    },
+    {
+      "epoch": 86.67,
+      "eval_loss": 4.51953125,
+      "eval_runtime": 2.8067,
+      "eval_samples_per_second": 3.919,
+      "eval_steps_per_second": 0.713,
+      "step": 260
+    },
+    {
+      "epoch": 90.0,
+      "learning_rate": 6e-06,
+      "loss": 0.0002,
+      "step": 270
+    },
+    {
+      "epoch": 93.33,
+      "learning_rate": 6e-06,
+      "loss": 0.0002,
+      "step": 280
+    },
+    {
+      "epoch": 93.33,
+      "eval_loss": 4.5390625,
+      "eval_runtime": 2.8074,
+      "eval_samples_per_second": 3.918,
+      "eval_steps_per_second": 0.712,
+      "step": 280
+    },
+    {
+      "epoch": 96.67,
+      "learning_rate": 6e-06,
+      "loss": 0.0002,
+      "step": 290
+    },
+    {
+      "epoch": 100.0,
+      "learning_rate": 6e-06,
+      "loss": 0.0003,
+      "step": 300
+    },
+    {
+      "epoch": 100.0,
+      "eval_loss": 4.5546875,
+      "eval_runtime": 2.8081,
+      "eval_samples_per_second": 3.917,
+      "eval_steps_per_second": 0.712,
+      "step": 300
+    },
+    {
+      "epoch": 100.0,
+      "step": 300,
+      "total_flos": 28875518115840.0,
+      "train_loss": 0.16138767729202907,
+      "train_runtime": 5518.2957,
+      "train_samples_per_second": 0.797,
+      "train_steps_per_second": 0.054
+    }
+  ],
+  "max_steps": 300,
+  "num_train_epochs": 100,
+  "total_flos": 28875518115840.0,
+  "trial_name": null,
+  "trial_params": null
+}

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e437e5065d5ecc3a58d100c2db2c4ab60408c0f8bf7a8837cff0a1ff7ef0598
+size 4923

version

@@ -0,0 +1 @@
+1.1.2