def model_fn()

in sagemaker_notebook_instance/containers/question_answering/entry_point.py [0:0]

• 16 lines of code
• 1 McCabe index (conditional complexity)

def model_fn(model_dir):
    session = sagemaker.Session()
    bucket = os.getenv("MODEL_ASSETS_S3_BUCKET")
    prefix = os.getenv("MODEL_ASSETS_S3_PREFIX")
    session.download_data(path=model_dir, bucket=bucket, key_prefix=prefix)
    tokenizer = AutoTokenizer.from_pretrained("bert-large-uncased-whole-word-masking-finetuned-squad", cache_dir=model_dir)
    model = AutoModelForQuestionAnswering.from_pretrained("bert-large-uncased-whole-word-masking-finetuned-squad", cache_dir=model_dir)
    answerer = pipeline(
        task="question-answering",
        model=model,
        tokenizer=tokenizer
    )
    model_assets = {
        "answerer": answerer
    }
    return model_assets