I am trying to just get embeddings, so the task I’m looking for is feature extraction I believe. However, I looked at Auto Classes but it looks like AutoFeatureExtractor is used for images? I get a not found error when I try this so not sure that is the right apporach eighter. In a nutshell here is what I’m trying to do.

Download a pretrained model for doing pooled embedding extractions
Fine tune this model on my own corpus (transactional data)
Run a batch embedding job on my data to extract the embedding features, that can then be used by a classifier down stream (which takes other inputs in addition to embeddings)

we are currently using sagemaker’s  BlazingText in this manner, but I wanted to give HF a try since we can use it for many different tasks (NER etc).
The reason I went with ‘facebook/bart-large-mnli’ was because I was able to successfully deploy this to an endpoint and get embeddings just like I want. In addition, it performed well on zero shot tasks in my domain, and in the deploy example on the model cart “feature-extraction” is listed as a task so I took that to mean it could be trained as such. However, when I deploy to an endpoint I get to specify “feature-extraction” as the task and it just works. Not sure how I can do something similar during finetuning, or if I need to look at a different model?
If look under ‘feature-extraction’ as a task in the model hub, the most popular model is “sentence-transformers/distilbert-base-nli-mean-tokens” . However, this model gives uniform probabilities for zero-shot (i.e. all predictions are ~ 1/num_classes), and it also gives me the same loss error when I try to fine tune so that’s why I didn’t use this. See below for the zero shot output for “sentence-transformers/distilbert-base-nli-mean-tokens” , if I use ‘facebook/bart-large-mnli’  instead I get expected more reasonable probabilities.

image1027×804 76.8 KB