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

# 사용자 입력 받기
user_input = input("텍스트 생성을 시작할 단어를 입력하세요: ")

# 텍스트 생성 함수 정의
def generate_text(seed_text, next_words, max_sequence_len):
    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 = model.predict_classes(token_list, verbose=0)
        
        output_word = ""
        for word, index in tokenizer.word_index.items():
            if index == predicted:
                output_word = word
                break
        seed_text += " " + output_word
    return seed_text

# 텍스트 생성
generated_text = generate_text(user_input, next_words=50, max_sequence_len=25)
print(generated_text)