from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForMaskedLM
from datasets import load_dataset

# 데이터셋 로드
dataset = load_dataset("HyeoniLEE/books_dataset")

# 데이터셋을 훈련 세트와 검증 세트로 나누기
dataset = dataset["train"].train_test_split(test_size=0.1)  # 10%를 검증 세트로 사용

# 토크나이저 및 모델 로드
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModelForMaskedLM.from_pretrained("bert-base-uncased")
)

def preprocess_function(examples):
    texts = []
    for i in range(len(examples["title"])):
        text = (
            (examples["title"][i] if examples["title"][i] is not None else "") + " " +
            (examples["author"][i] if examples["author"][i] is not None else "") + " " +
            (examples["table_of_contents"][i] if examples["table_of_contents"][i] is not None else "") + " " +
            (examples["book_intro"][i] if examples["book_intro"][i] is not None else "") + " " +
            (examples["publisher_review"][i] if examples["publisher_review"][i] is not None else "") + " " +
            (examples["review"][i] if examples["review"][i] is not None else "")
        )
        texts.append(text)
    tokenized_inputs = tokenizer(texts, truncation=True, padding="max_length", max_length=512)
    tokenized_inputs["labels"] = tokenized_inputs["input_ids"].copy()  # labels 추가
    return tokenized_inputs

# 데이터셋 전처리
tokenized_dataset = dataset.map(preprocess_function, batched=True)

training_args = TrainingArguments(
    output_dir="./results",
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    num_train_epochs=3,
    weight_decay=0.01,
)

dataset = dataset["train"].train_test_split(test_size=0.1)

10%를 검증 데이터셋으로 사용