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()