import tensorflow as tf
from tensorflow.keras.preprocessing.text import Tokenizer
from tensorflow.keras.preprocessing.sequence import pad_sequences

# 문서 데이터
documents = [
    "이 휴대폰은 화면이 크고 배터리 수명도 길어요.",
    "이 노트북은 성능이 좋고 키보드도 편해요.",
    "이 스피커는 음질이 좋고 크기도 작아서 휴대하기 편해요.",
    "이 카메라는 사진 화질이 좋고 조작이 간단해요.",
    "이 마우스는 디자인이 멋지고 사용하기 편해요."
]

# 레이블 데이터
labels = [
    "휴대폰",
    "노트북",
    "스피커",
    "카메라",
    "마우스"
]

# 텍스트를 숫자로 변환
tokenizer = Tokenizer()
tokenizer.fit_on_texts(documents)
sequences = tokenizer.texts_to_sequences(documents)

# 문장의 길이를 맞춰줌
max_len = max(len(seq) for seq in sequences)
padded_sequences = pad_sequences(sequences, maxlen=max_len, padding='post')

# 레이블을 숫자로 변환
label_tokenizer = Tokenizer()
label_tokenizer.fit_on_texts(labels)
encoded_labels = label_tokenizer.texts_to_sequences(labels)

# 모델 구축
model = tf.keras.Sequential([
    tf.keras.layers.Embedding(len(tokenizer.word_index)+1, 16, input_length=max_len),
    tf.keras.layers.GlobalAveragePooling1D(),
    tf.keras.layers.Dense(16, activation='relu'),
    tf.keras.layers.Dense(5, activation='softmax')
])

# 모델 컴파일
model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

# 모델 훈련
model.fit(padded_sequences, 
          encoded_labels, 
          epochs=10, 
          verbose=2)

# 문서 분류 예측
new_documents = [
  "이 키보드는 무게가 가볍고 조작이 편리해요.",
  "이 스마트폰은 카메라가 우수하고 배터리도 오래가요."
]

new_sequences = tokenizer.texts_to_sequences(new_documents)
new_padded_sequences = pad_sequences(new_sequences, maxlen=max_len, padding='post')
predictions = model.predict(new_padded_sequences)

for i, document in enumerate(new_documents):
    predicted_label = label_tokenizer.sequences_to_texts([tf.argmax(predictions[i]).numpy()])[0]
    print(f"{document} --> {predicted_label}")