# Workaround to install the lib without "setup.py"
import sys
from git import Repo
Repo.clone_from("https://github.com/dimitreOliveira/hub.git", "./hub")
sys.path.append("/hub")

import gradio as gr
import tensorflow as tf
import tensorflow_text as text  # Registers the ops. (needed for "bert_preprocessor")
from hub.tensorflow_hub.hf_utils import pull_from_hub


def model_fn(preprocessor, encoder):
    text_input = tf.keras.layers.Input(shape=(), dtype=tf.string)
    encoder_inputs = preprocessor(text_input)
    outputs = encoder(encoder_inputs)
    pooled_output = outputs["pooled_output"]
    return tf.keras.Model(text_input, pooled_output)

def predict_fn(text):
    print(text)
    print(tf.constant([text]))
    embed = model(tf.constant([text]))[0].numpy()
    return embed

preprocessor = pull_from_hub(repo_id="Dimitre/bert_en_cased_preprocess")
encoder = pull_from_hub(repo_id="Dimitre/bert_en_cased_L-12_H-768_A-12")
model = model_fn(preprocessor, encoder)

iface = gr.Interface(fn=predict_fn, 
                     title="BERT sentence embeddings", 
                     description="Get the embeddings from your sentences using BERT", 
                     inputs=gr.Textbox(lines=2, placeholder="Text input here...", label="Text"), 
                     outputs="text", 
                     examples=[["Hello! This is a random sentence"]])
iface.launch()