import tensorflow as tf
import gradio as gr
import pickle
import os

# define variables
checkpoint_folder = './'
bi_gru_name = "test_25_128_2564_15_32_10_16_05_D_64_15.pickle"
tokenizer_path = os.path.join(checkpoint_folder, 'tokenizer_tP4suwZ.pickle')
model_path = os.path.join(checkpoint_folder, bi_gru_name)
INDEX_CLASS = {0:'suicide', 1:'non-suicide'}
MAX_SEQ_DF = 64


# load the tokenizer from pickle
with open(tokenizer_path, 'rb') as tokenizer_file:
    tokenizer = pickle.load(tokenizer_file)
# load the mode from pickle and read the model
with open(model_path, 'rb') as model_file:
    model_json = pickle.load(model_file)
model = tf.keras.models.model_from_json(model_json)

# define prediction function
def predict(input_text):
    # preprocessing
    tokenized_data = tokenizer.texts_to_sequences([input_text])
    text_data_padded = tf.keras.preprocessing.sequence.pad_sequences(tokenized_data, maxlen = MAX_SEQ_DF, padding = 'post')

    # make prediction
    pred = model.predict(text_data_padded)
    prediction = INDEX_CLASS[round(pred[0][0])]

    return prediction


gr_intrfc = gr.Interface(fn=predict, inputs="text", title='Suicide Detection',
            examples=['I Just want it to stop, what the point anyway',
                      'Thanks for the help, i do nont feel depressed anymore'],
                      outputs="text", theme='dark')
gr_intrfc.launch(debug=True, )