TypeError: model_fn() takes 1 positional argument but 2 were given
I am trying to deploy Sagemaker Endpoint with custom scripts using model_fn, predict fn and input_fn using sqlcoder7b-2 and I am getting errorTypeError: model_fn() takes 1 positional argument but 2 were given#code/inference.py
def model_fn(model_dir):
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    model = AutoModelForCausalLM.from_pretrained(
        model_dir,
        trust_remote_code=True,
        torch_dtype=torch.float16,
        device_map="auto",
        use_cache=True,
    cache_dir ="sqlcoder",
    offload_folder="offload_sqlcoder",
    force_download=True
    )
    model_dict = {'model':model, 'tokenizer':tokenizer}
    return model_dict 

def predict_fn(data, model_dict):
    tokenizer = model_dict['tokenizer']
    model = model_dict['model']
    question=data["inputs"]
    print("Question is",question)
    print("Prompt is",prompt)
    updated_prompt = prompt.format(question=question)
    print("Updated Prompt is",prompt)
    inputs = tokenizer(updated_prompt, return_tensors="pt").to("cuda")
    generated_ids = model.generate(
        **inputs,
        num_return_sequences=1,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.eos_token_id,
        max_new_tokens=400,
        do_sample=False,
        num_beams=1,
    )
    outputs = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

    torch.cuda.empty_cache()
    torch.cuda.synchronize()
    return sqlparse.format(outputs[0].split("[SQL]")[-1], reindent=True)

    
def input_fn(request_body, request_content_type):
    # Transform the input request to a dictionary
    print("Request body is",request_body)
    print("Request Content Type is", request_content_type)
    request = json.loads(request_body)
    
    return request#Deployment Script
from sagemaker.huggingface.model import HuggingFaceModel

hub = {
#  'HF_MODEL_ID':'defog/sqlcoder-7b-2', # model_id from hf.co/models
  'HF_TASK':'text-generation' # NLP task you want to use for predictions
}
# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
   model_data=s3_location,       # path to your model and script
   role=role,                    # iam role with permissions to create an Endpoint
   transformers_version="4.37.0",  # transformers version used
   pytorch_version="2.1.0",        # pytorch version used
   py_version='py310',            # python version used
   env=hub
)

# deploy the endpoint endpoint
predictor = huggingface_model.deploy(
    initial_instance_count=1,
    instance_type="ml.g5.xlarge"
    )
2024-03-13T00:26:44Z
Did you fix the issue?
2024-03-13T18:39:23Z
Yes, I was able to resolve the issue. I removed my function “input_fn” from the script (this meant that I ended up using “input” func written in sagemaker huggingface inference toolkit). In “model_fn” and “predict_fn”, I added an argument “context=None”. In predict_fn, I modified argument “model_dict” to “model”.
2024-04-02T13:53:48Z
How did you know to do this? I’m currently facing this issue having just redeployed a previously working endpoint + endpoint configuration. Not sure what the issue is.
2024-04-02T14:19:28Z
This code is executed for inference. You can see lines 277 to 281 where they have defined functions load_fn, predict_fn etc. These functions are like user defined proxies to load, predict and other huggingface functions. Just one constraint that these user proxy functions should take the arguments passed by respective huggingface defined functions.github.comaws/sagemaker-huggingface-inference-toolkit/blob/main/src/sagemaker_huggingface_inference_toolkit/handler_service.py# Copyright 2021 The HuggingFace Team, Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import importlib
import logging
import os
import sys
import time
from abc import ABCThis file has been truncated.show original
2024-04-04T14:30:40Z
This topic was automatically closed 12 hours after the last reply. New replies are no longer allowed.