import gradio as gr
import tensorflow as tf
from SpaceGen_preprocessing import *
from utils import *


# Path to your Keras model
model_path = "SpaceGen_Large.keras"

# Load the model
model = tf.keras.models.load_model(model_path)


def fix_space(text):
  text = clean_sentence(text)
  X = text_to_X(text)
  predictions = model.predict(X, verbose=0)
  predicted_labels = []
  for pred in predictions[0]:
    predicted_labels.append(1 if pred[1] > .5 else 0)
  fixed_text = insert_spaces(text.replace(' ',''), find_indices(predicted_labels))
  return fixed_text

default_text = "T hel ittlegi rlra nthro ughth epa rkc has ing abut terfly."
demo = gr.Interface(fn=fix_space,
                    inputs=gr.Textbox(label="Input Text", value=default_text),
                    outputs="text")
demo.launch()