Inference with My Own Embedding Model

Download original notebook: [own embedding model inference](/notebooks/own embedding model inference)

!pip install safetensors

from safetensors import safe_open

tensors = {}
with safe_open("/content/drive/MyDrive/embedding_model/model.safetensors", framework="pt", device=0) as f:
    for k in f.keys():
        tensors[k] = f.get_tensor(k)

!pip install accelerate==0.24.1
!pip install huggingface-hub==0.17.3
!pip install safetensors==0.4.0
!pip install tokenizers==0.14.1
!pip install transformers==4.35.0

model = AutoModelForMaskedLM.from_pretrained("/content/drive/MyDrive/embedding_model/model.safetensors")

from transformers import AutoTokenizer, AutoModel
import torch


def cls_pooling(model_output, attention_mask):
    return model_output[0][:,0]


# Sentences we want sentence embeddings for
sentences = ['This is an example sentence', 'Each sentence is converted']

# Load model from HuggingFace Hub
tokenizer = AutoTokenizer.from_pretrained('shrijayan/embedding_train')
model = AutoModel.from_pretrained('shrijayan/embedding_train')

# Tokenize sentences
encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')

# Compute token embeddings
with torch.no_grad():
    model_output = model(**encoded_input)

# Perform pooling. In this case, cls pooling.
sentence_embeddings = cls_pooling(model_output, encoded_input['attention_mask'])

print("Sentence embeddings:")
print(sentence_embeddings)

len(sentence_embeddings[0])