Cannot pass NeuronModelForCausalLM to pipeline

[philschmid]

I cannot pass a model initiated from NeuonModelForCausalLM into the pipeline directly. See error below

    model = NeuronModelForCausalLM.from_pretrained(model_dir, low_cpu_mem_usage=True)
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    print("loading pipeline...")
    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

creates

 Model <optimum.neuron.modeling.NeuronModelForCausalLM object at 0x7fa4f8131d20> is not supported. Please provide a valid model either as string or NeuronModel. You can also provide non model then a default one will be used
Traceback (most recent call last):
File "/opt/conda/lib/python3.10/site-packages/mms/model_loader.py", line 145, in load  model_service.initialize(service.context)
File "/opt/conda/lib/python3.10/site-packages/sagemaker_huggingface_inference_toolkit/handler_service.py", line 77, in initialize
self.model = self.load(self.model_dir)
File "/opt/ml/model/code/inference.py", line 18, in model_fn
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
File "/opt/conda/lib/python3.10/site-packages/optimum/neuron/pipelines/transformers/base.py", line 229, in pipeline
model, model_id, tokenizer, feature_extractor = load_pipeline(
File "/opt/conda/lib/python3.10/site-packages/optimum/neuron/pipelines/transformers/base.py", line 148, in load_pipeline
raise ValueError(
 ValueError: Model <optimum.neuron.modeling.NeuronModelForCausalLM object at 0x7fa4f8131d20> is not supported. Please provide a valid model either as string or NeuronModel.
You can also provide non model then a default one will be used

[dacorvo]

Good catch ! As a workaround until I fix this, you can still do:

tokenizer = AutoTokenizer.from_pretrained(model_dir)
print("loading pipeline...")
pipe = pipeline("text-generation", model_dir, tokenizer=tokenizer, low_cpu_mem_usage=True)

[dacorvo]

Here is the fix: #319