BERT를 이용한 문장 관계 예측

!pip install transformers kobert-transformers Korpora

from Korpora import KoreanChatbotKorpus

data = KoreanChatbotKorpus()

data.train[0].text

data.train[1].pair

from kobert_transformers import get_tokenizer

tokenizer = get_tokenizer()

from transformers import TFBertForNextSentencePrediction

model = TFBertForNextSentencePrediction.from_pretrained('monologg/kobert', from_pt=True)

import random

from sklearn.model_selection import train_test_split

train_data, test_data = train_test_split(data.train, test_size=0.2, random_state=1234)

train_texts = []
train_pairs = []
train_labels = []
for item in train_data:
    train_texts.append(item.text)
    label = random.randint(0, 1)
    train_labels.append(label)
    if label == 1:
        train_pairs.append(item.pair)
    else:
        fake = random.choice(train_data)
        train_pairs.append(fake.pair)

test_texts = []
test_pairs = []
test_labels = []
for item in test_data:
    test_texts.append(item.text)
    label = random.randint(0, 1)
    test_labels.append(label)
    if label == 1:
        test_pairs.append(item.pair)
    else:
        fake = random.choice(test_data)
        test_pairs.append(fake.pair)

train_encodings = tokenizer(train_texts, train_pairs, padding=True, truncation=True)
test_encodings = tokenizer(test_texts, test_pairs, padding=True, truncation=True)

train_encodings.keys()

train_encodings['input_ids'][0]

train_encodings['token_type_ids'][0]

train_encodings['attention_mask'][0]

import tensorflow as tf

train_dataset = tf.data.Dataset.from_tensor_slices((
    dict(train_encodings),
    train_labels
))

test_dataset = tf.data.Dataset.from_tensor_slices((
    dict(test_encodings),
    test_labels
))

optimizer = tf.keras.optimizers.Adam(learning_rate=5e-5)
model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'])

model.evaluate(test_dataset.batch(32), batch_size=32)

model.fit(train_dataset.batch(32), epochs=1, batch_size=32)

model.evaluate(test_dataset.batch(32), batch_size=32)

pred_dataset = tf.data.Dataset.from_tensor_slices(dict(test_encodings))

prediction = model.predict(pred_dataset.batch(32))

tf.math.argmax(prediction[0], axis=-1).numpy()