transformers_inference

trans_inference.py

# Sentences we want sentence embeddings for
sentences = ['This is an example sentence', 'This is sample of the sentence']



import time
start = time.time()



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
sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])

# Normalize embeddings
sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
#calculate similarity
cosine_scores = util.pytorch_cos_sim(sentence_embeddings[0], sentence_embeddings[1])
cosine_scores


end = time.time()

print(end - start)
print(f"pytorch vanilla cpu: {(end- start)/2:.2f}s/sequence")