embed.py

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]}