Unable to deploy fine tuned Mistral
I am trying to deploy a fine tunedmistralai/Mistral-7B-v0.1model to Sagemaker. Along with the base model, we also have an adapter to load. The adapter files and model files are packaged into a.tar.gzfile.Here is ourinference.pyimport torch
import json
from transformers import AutoTokenizer, AutoModel
from torch import nn

class AutoModelForSentenceEmbedding(nn.Module):
    def __init__(self, base_model, tokenizer, normalize=True):
        super(AutoModelForSentenceEmbedding, self).__init__()

        self.model = base_model  # , load_in_8bit=True, device_map={"":0})
        self.normalize = normalize
        self.tokenizer = tokenizer

    def forward(self, **kwargs):
        ea = self.get_embeddings(kwargs['input_ids'], kwargs['attention_mask'])
        ep = self.get_embeddings(kwargs['input_ids_positive'], kwargs['attention_mask_positive'])
        en = self.get_embeddings(kwargs['input_ids_negative'], kwargs['attention_mask_negative'])
        return ea, ep, en
    
    def get_embeddings(self, input_ids, attention_mask):
        model_output = self.model(input_ids, attention_mask)
        embeddings = self.mean_pooling(model_output, attention_mask)
        if self.normalize:
            embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
        return embeddings

    def mean_pooling(self, model_output, attention_mask):
        token_embeddings = model_output[0]  # First element of model_output contains all token embeddings
        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

    def __getattr__(self, name: str):
        """Forward missing attributes to the wrapped module."""
        try:
            return super().__getattr__(name)  # defer to nn.Module's logic
        except AttributeError:
            return getattr(self.model, name)
        
def model_fn(model_dir, context=None):
    model = AutoModel.from_pretrained(model_dir, device_map='auto', load_in_8bit=True)
    tokenizer = AutoTokenizer.from_pretrained(model_dir)

    model.load_adapter(f"{model_dir}/adapter")
    print("loaded adapter")
    return AutoModelForSentenceEmbedding(model, tokenizer)

def transform_fn(model, input_data, content_type, accept):
     # decode the input data (e.g. JSON string -> dict)
    data = decoder_encoder.decode(input_data, content_type)

    # call your custom model with the data
    outputs = model.get_embeddings("blast")

    # convert the model output to the desired output format (e.g. dict -> JSON string)
    response = decoder_encoder.encode(output, accept)

    return responserequirements.pyaccelerate
transformers
bitsandbytesAnd our notebook hasfrom sagemaker.huggingface.model import HuggingFaceModel

huggingface_model = HuggingFaceModel(
   model_data=s3_model_uri,      # path to your model and script
   role=role,                    # iam role with permissions to create an Endpoint
   transformers_version="4.26",                           # Transformers version used
   pytorch_version="1.13",                                # PyTorch version used
   py_version='py39'
)

predictor = huggingface_model.deploy(
   initial_instance_count=1,
   instance_type='ml.g4dn.xlarge',
    endpoint_name='query-embeddings'
)We are encountering the following error when we are calling the predict.com.amazonaws.ml.mms.wlm.WorkerLifeCycle - KeyError: 'mistral'From the looks of it, it looks like a transformer version mismatch but we are not able to figure out what is the correct version to go for. Has someone deployed a fine tuned mistral model before? If so, what is the correct approach?