# Extract hidden layers from a Roberta model in sagemaker

**URL:** </static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fextract-hidden-layers-from-a-roberta-model-in-sagemaker%2F10748%3E%5C
**Category:** Beginners\
**Created:** [October 13, 2021, 5:37pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fextract-hidden-layers-from-a-roberta-model-in-sagemaker%2F10748 "2021-10-13T17:37:17Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![StanBienaives](https://avatars.discourse-cdn.com/v4/letter/s/858c86/32.png) [@StanBienaives](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FStanBienaives)%5C
**Post date:** [October 13, 2021, 5:37pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fextract-hidden-layers-from-a-roberta-model-in-sagemaker%2F10748%2F1 "2021-10-13T17:37:18Z")

</div>

Hello,

I have fine tuned a Camembert Model (inherits from Roberta) on a custom dataset using sagemaker.  
My goal is to have a language model able to extract embedding to be used in my search engine.

Camembert is trained for a “fill-mask” task.  
Using the Huggingface API outputting hidden\_layers (thus computing embedding) is fairly simple  
`model = AutoModelForMaskedLM.from_pretrained(args.model_name, output_hidden_states=True)`

But when deploying such model in sagemaker the predict method only returns the text output.  
There is some kind of post-processing that I do not control.

Is there a way to customize the post-processing steps in sagemaker ?  
What model architecture should I be using to extract embeddings on sagemaker ?

Thanks for your help.

---

<div class="post-metadata">

**Author:** ![StanBienaives](https://avatars.discourse-cdn.com/v4/letter/s/858c86/32.png) [@StanBienaives](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FStanBienaives)%5C
**Post date:** [October 21, 2021, 10:20am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fextract-hidden-layers-from-a-roberta-model-in-sagemaker%2F10748%2F2 "2021-10-21T10:20:43Z")

</div>

For those having the same issue I found a solution.

Train the model on masked ML and at inference time use the pipeline ‘feature\_extraction’ by setting the HF\_TASK environment variable.

```auto
hub = {
  'HF_TASK': 'feature_extraction' 
}
huggingface_model = HuggingFaceModel(
   env=hub,
   model_data="s3://bucket/model.tar.gz"
   role=<SageMaker Role>, 
   transformers_version="4.6", 
   pytorch_version="1.7", 
   py_version="py36",
)
huggingface_model.deploy(1, '<ec2 -instance type >')

```

the model will send the feature vector for each token. To get the sentence vector you can average the word vectors or use other fancy methods I dit not explore.

If you want more control on what exactly your model returns you can customize the  
output\_fn, predict\_fn etc… as described here: [GitHub - aws/sagemaker-huggingface-inference-toolkit](https://github.com/aws/sagemaker-huggingface-inference-toolkit#-user-defined-codemodules)