Ubuntu
.
270cf00
raw
history blame contribute delete
816 Bytes
import gradio as gr
import json
from transformers import AutoTokenizer, TFAutoModel
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_pooling(model_output):
return model_output.last_hidden_state[:, 0]
def get_embeddings(text_list):
encoded_input = tokenizer(
text_list, padding=True, truncation=True, return_tensors="tf"
)
encoded_input = {k: v for k, v in encoded_input.items()}
model_output = model(**encoded_input)
return json.dumps(cls_pooling(model_output).numpy().tolist())
gradio_interface = gr.Interface(
fn = get_embeddings,
inputs = gr.Textbox(lines=10, placeholder="Input Text"),
outputs = "text"
)
gradio_interface.launch()