raw_datasets = datasets.load_dataset("imdb")

#사전 학습모델선택, 토크나이저 로드
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)

#토크나이징, 전처리
def tokenize_function(examples):
    return tokenizer(
        examples["text"], truncation=True, padding="max_length", max_length=512)
 
#모든 샘플에 토크나이징 함수를 mapping하기       
tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)

#데이터셋 전처리 및 준비
tokenized_datasets = tokenized_datasets.remove_columns(["text"])
tokenized_datasets.set_format("torch")

# 최종 학습 및 검증 데이터셋 준비
train_dataset = tokenized_datasets["train"]
eval_dataset = tokenized_datasets["test"]

# 3. 모델 로드
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)

import numpy as np
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score

# 4. 평가 지표 정의
def compute_metrics(eval_pred):
    # 평가 예측 결과를 사용하여 정확도, 정밀도, 재현율, F1 점수를 계산합니다.
    predictions, labels = eval_pred
    # 가장 logit이 높은값만 1로  preds에 저장한다.
    preds = np.argmax(predictions, axis = 1)
    accuracy = accuracy_score(labels, preds)
    precision = precision_score(labels, preds)
    recall = recall_score(labels, preds)
    f1 = f1_score(labels, preds)
    return {
        "accuracy": accuracy,
        "precision": precision,
        "recall": recall,
        "f1": f1,
    }

# 5. 학습 설정
# TrainingArguments를 사용하여 학습 설정을 정의합니다.
# TrainingArguments는 모델 학습에 필요한 다양한 하이퍼파라미터를 설정합니다.
# 여러 변수를 필요에 따라 조정해보세요.
# 예:
# 학습 결과(모델 체크포인트 등), 로그를 저장할 디렉토리
# 모델을 평가나 저장할 방식
# 학습률 설정 (파라미터 업데이트 크기)
# 디바이스(GPU/CPU) 당 학습 배치 크기
# 디바이스 당 평가 배치 크기
# 전체 학습 데이터를 몇 번 반복할지 (에폭 수)
# 과적합 방지를 위한 가중치 감쇠 계수
# 로그를 기록할 스텝 간격
training_args = TrainingArguments(output_dir = "./results",
                                  num_train_epochs = 2, # 총 훈련 에폭 수
                                  per_device_train_batch_size = 16, #device당 훈련 배치 크기
                                  per_device_eval_batch_size = 64, #device당 평가 배치 크기
                                  learning_rate = 5e-5,
                                  weight_decay = 0.01,

                                  # evaluation_strategy = "epoch", #epoch마다 평가
                                  # save_strategy = "epoch", #epoch마다 모델저장
                                  logging_dir = "./logs",
                                  logging_steps = 100, #100스탭마다 기록
                                  # load_best_model_at_end = True #훈련 종료시 가장 성능이 좋았던 모델을 로드
                                  )

# 6. Trainer 정의 및 학습
# Trainer를 사용하여 모델을 학습하고 평가합니다.
# Trainer는 모델, 학습 인자, 데이터셋 등을 포함하여 학습 및 평가를 관리합니다.
# training_args 사용하여 학습 설정을 정의합니다.
# 앞서 정의한 compute_metrics 함수를 사용하여 평가 지표를 계산합니다.
trainer = Trainer(model = model,
                  args = training_args,
                  train_dataset = train_dataset,
                  eval_dataset = eval_dataset,
                  compute_metrics = compute_metrics)

trainer.train()
trainer.evaluate()

# 7. 새 문장 예측
model.eval()  # 모델을 평가 모드로 설정
new_sentences = [
    "This movie was absolutely fantastic! I highly recommend it.",
    "What a terrible film. I wasted my time watching it.",
    "The plot was a bit confusing, but the acting was superb.",
    "I loved every single moment of this incredible story.",
]

inputs = tokenizer(
    new_sentences, return_tensors="pt", padding=True, truncation=True, max_length=512
)
inputs = {k: v.to(model.device) for k, v in inputs.items()}
with torch.no_grad():
    outputs = model(**inputs)
    probs = torch.softmax(outputs.logits, dim=-1)

for i, sentence in enumerate(new_sentences):
    pred = torch.argmax(probs[i]).item()
    label = "긍정 (Positive)" if pred == 1 else "부정 (Negative)"
    print(f"\n문장: '{sentence}'")
    print(f"예측: {label} (부정 확률: {probs[i][0]:.4f}, 긍정 확률: {probs[i][1]:.4f})")

with torch.no_grad():
    outputs = model(**inputs)
    print(outputs.logits)