embed.py

from transformers import AutoTokenizer,TFAutoModel
import torch
from torch import nn
import tensorflow



model_ckpt = "sentence-transformers/multi-qa-mpnet-base-dot-v1"

tokenizer=AutoTokenizer.from_pretrained(model_ckpt)

model=TFAutoModel.from_pretrained(model_ckpt,from_pt=True)


def cls_pool(model):

  return model.last_hidden_state[:,0,:]

def sample_embedding(example):

  token_output=tokenizer(example,padding=True,truncation=True,return_tensors="tf")

  token_output={k:v for k,v in token_output.items()}
  

  model_output=model(**token_output)

  return {"embedding":cls_pool(model_output).numpy()[0]}