import gradio as gr
import tensorflow as tf
from tensorflow.keras.preprocessing.sequence import pad_sequences
import numpy as np
import pickle

# 토크나이저 불러오기
with open('tokenizer.pickle', 'rb') as handle:
    tokenizer = pickle.load(handle)

# 텍스트 생성 모델 불러오기
model = tf.keras.models.load_model("text_generation_model")

# 텍스트 생성 함수 정의
def generate_text(seed_text, next_words=50, max_sequence_len=25):
    for _ in range(next_words):
        token_list = tokenizer.texts_to_sequences([seed_text])[0]
        token_list = pad_sequences([token_list], maxlen=max_sequence_len-1, padding='pre')
        predicted = np.argmax(model.predict(token_list), axis=-1)
        
        output_word = ""
        for word, index in tokenizer.word_index.items():
            if index == predicted:
                output_word = word
                break
        seed_text += " " + output_word
    return seed_text

# Gradio 인터페이스 정의
iface = gr.Interface(
    fn=generate_text,
    inputs=[gr.Textbox(label="시작할 단어를 입력하세요", placeholder="여기에 입력하세요..."), gr.Slider(label="생성할 단어 수", min_value=1, max_value=100, default=50), gr.Slider(label="최대 시퀀스 길이", min_value=1, max_value=100, default=25)],
    outputs=gr.Textbox(label="생성된 텍스트"),
    title="텍스트 생성기",
    description="사용자가 입력한 텍스트를 기반으로 추가 텍스트를 생성합니다."
)

iface.launch()