How does the predict function know what model to use?

[la-cruche]

Hi,

In current SageMaker Framework hosting (eg Sklearn-on-Flask, PyTorch-on-Torchserve), the prediction function is a predict(input_object, model) ; that takes as input the result of the model_fn.

In SM HF Hosting the predict only receives the processed_data. Then how does it know on which model to work? the objects returned by load_fn are available in memory?

[philschmid]

There is a PR for this already, which will be merged today. #2
This PR adds the model as input parameter for predict

[la-cruche]

ok perfect. Yes indeed, with the current predict_fn(processed_data) the inference fails with a W-model-1-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - mms.service.PredictionException: name 'model' is not defined : 400

I'm using the below functions

def load_fn(model_dir):
    """this function reads the model from disk"""
    
    print('load_fn dir view:')
    print(os.listdir())
    
    # load model
    model = TFAutoModelForQuestionAnswering.from_pretrained('/opt/ml/model')
    
    # load tokenizer
    tokenizer = AutoTokenizer.from_pretrained('/opt/ml/model')
    
    return model, tokenizer




def predict_fn(processed_data):
    """this function runs inference"""

    print('processed_data received: ')
    print(processed_data)
    
    print('model name:')
    print(model.name)

    print('tok name:')
    print(tokenizer.name)   
    
    question, text = processed_data['inputs']['question'], processed_data['inputs']['context']

    input_dict = tokenizer(question, text, return_tensors='tf')
    outputs = model(input_dict)
    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

[la-cruche]

with the new predict_fn(processed_data, model) things work fine as documented here #10 (comment)