Ability to transform model outputs in DJL Serving

[rachitchauhan43]

Description

We are using SageMaker for large model inference (LMI) as documented here

With this notebook https://github.com/deepjavalibrary/djl-demo/blob/master/aws/sagemaker/large-model-inference/sample-llm/rollingbatch_llama_7b_customized_preprocessing.ipynb, we saw there is a way to manipulate the input before it goes to then model because parse_input method is available as a hook but we also have a need to manipulate the output before it leaves model server.

@lanking520 and @frankfliu Any thoughts on supporting that ?

Will this change the current api? How? No. It's just an extension we are asking for.

Benefit of this will be: Users won;t have to write another service layer in front of model server just to manipulate/transform outputs.

Who will benefit from this enhancement?

References

• https://github.com/deepjavalibrary/djl-demo/blob/master/aws/sagemaker/large-model-inference/sample-llm/rollingbatch_llama_7b_customized_preprocessing.ipynb
• https://docs.aws.amazon.com/sagemaker/latest/dg/large-model-inference-dlc.html

[rachitchauhan43]

cc: @chirag-orbittec

[frankfliu]

@rachitchauhan43

We do have plan to allow customize output. We have a few built-in output_formatter for rolling batch. We are considering allows user to point the output_formatter to a user provided module.

Are you considering to contribute to this feature?

[chirag-orbittec]

@frankfliu I can contribute that feature. We need this