inference.py ignored by SM HF about sagemaker-huggingface-inference-toolkit

[philschmid]

It is ignored because you define load in your inference.py and not as documented here load_fn here. That's why the toolkit uses the load function of the HuggingFaceHandlerService.
Additionally, I think the model in your predict might not be used since currently, the HuggingFaceHandlerService uses self.model in the predict function.

[la-cruche]

ok I'm lost :)

how should we name the functions in inference.py? (also asked here #6)

1. load_fn, preprocess_fn, predict, postprocess (what's currently in the doc)
2. load_fn, preprocess_fn, predict_fn, postprocess_fn

I understand that it will soon switch to model_fn, input_fn, predict_fn, output_fn ; but I'm curious about right now

[philschmid]

Here is the example we use for the tests. Yes, the documentation has an issue. It should be _fn everywhere.

https://github.com/aws/sagemaker-huggingface-inference-toolkit/blob/main/tests/resources/code/inference.py

and here is it where they got loaded

sagemaker-huggingface-inference-toolkit/src/sagemaker_huggingface_inference_toolkit/handler_service.py

Line 209 in a3e85a1

def validate_and_initialize_user_module(self):

[la-cruche]

this is the final setup that worked ; as mentioned in the doc the contract now is model_fn, input_fn, predict_fn, output_fn

import logging
import os

import tensorflow as tf
from transformers import TFAutoModelForQuestionAnswering, AutoTokenizer


logging.basicConfig(level=logging.INFO)



def model_fn(model_dir):
    """this function reads the model from disk"""
    
    logging.info('model_fn dir view:')
    logging.info(os.listdir())
        
    # load model
    transformer = TFAutoModelForQuestionAnswering.from_pretrained(model_dir)
    
    # load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
        
    return transformer, tokenizer



def predict_fn(processed_data, model):
    """this function runs inference"""
        
    transformer, tokenizer = model
    question, text = processed_data['question'], processed_data['context']

    logging.info('processed_data received: {}'.format(processed_data))
    
    # infer
    input_dict = tokenizer(question, text, return_tensors='tf')
    outputs = transformer(input_dict)
    
    # post processing
    start_logits = outputs.start_logits
    end_logits = outputs.end_logits
    all_tokens = tokenizer.convert_ids_to_tokens(input_dict["input_ids"].numpy()[0])
    answer = ' '.join(all_tokens[tf.math.argmax(start_logits, 1)[0] : tf.math.argmax(end_logits, 1)[0]+1])
    
    return answer