import torch
import numpy as np
from datasets import load_dataset, Dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer
from sklearn.metrics import accuracy_score, f1_score

# 1. GPU/CPU 장치 설정 (학습 시 Trainer가 자동으로 처리하므로 확인용)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"사용 장치: {device}")

# 2. 모델 및 토크나이저 로드
MODEL_NAME = "FacebookAI/xlm-roberta-base"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
# 분류를 위해 AutoModelForSequenceClassification 로드 (분류 헤드가 추가됨)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2) 
# num_labels=2: 긍정(1), 부정(0)을 구분하도록 설정

# 3. 가상 데이터셋 준비 및 전처리
# 실제 학습 시에는 여기에 당신의 당근마켓 리뷰 데이터를 로드해야 합니다.
data = {
    'text': [
        "매너가 정말 좋으세요! 기분 좋은 거래였습니다.",  # 긍정
        "물건 상태가 별로고 답변도 너무 느렸어요.",      # 부정
        "빠른 응답과 깔끔한 거래 감사합니다.",          # 긍정
        "가격이 너무 비싸네요. 비추입니다.",          # 부정
        "오늘도 만족스러운 중고 거래였습니다.",         # 긍정
        "시간 약속 안 지키고 연락도 잘 안 되네요.",     # 부정
    ],
    'label': [1, 0, 1, 0, 1, 0] # 1: 긍정, 0: 부정
}
raw_dataset = Dataset.from_dict(data)

# 데이터셋을 학습(train)과 평가(test) 세트로 분할 (예시이므로 50:50)
train_test_split = raw_dataset.train_test_split(test_size=0.5, seed=42)
train_dataset = train_test_split['train']
eval_dataset = train_test_split['test']

def tokenize_function(examples):
    # 입력 텍스트를 토큰화합니다.
    return tokenizer(examples['text'], truncation=True, padding='max_length', max_length=128)

# 데이터셋에 토크나이저 적용
tokenized_train_dataset = train_dataset.map(tokenize_function, batched=True)
tokenized_eval_dataset = eval_dataset.map(tokenize_function, batched=True)

# 4. 평가 지표 함수 정의
def compute_metrics(p):
    # 예측된 로짓(logits)에서 argmax를 취해 예측 레이블을 얻습니다.
    predictions = np.argmax(p.predictions, axis=1)
    # 정확도(Accuracy)와 F1-Score를 계산합니다.
    acc = accuracy_score(p.label_ids, predictions)
    f1 = f1_score(p.label_ids, predictions, average='binary') # 긍정(1)에 대한 F1-Score
    return {"accuracy": acc, "f1": f1}

# 5. 학습 설정 (TrainingArguments)
OUTPUT_DIR = "./xlm-roberta-review-classifier" # 모델을 저장할 경로
training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    num_train_epochs=3,                     # 학습 횟수 (실제 작업 시 3~5회 권장)
    per_device_train_batch_size=8,          # GPU당 학습 배치 크기 (VRAM에 따라 조정)
    per_device_eval_batch_size=8,           # GPU당 평가 배치 크기
    warmup_steps=500,                       # 학습률이 최대치에 도달하는 단계 수
    weight_decay=0.01,                      # 가중치 감소 (오버피팅 방지)
    logging_dir='./logs',                   # 로그 저장 경로
    logging_steps=10,
    eval_strategy="epoch",            # 에포크마다 평가 수행
    save_strategy="epoch",                  # 에포크마다 모델 저장
    load_best_model_at_end=True,            # 학습 종료 시 가장 좋은 성능의 모델 로드
    fp16=torch.cuda.is_available(),         # GPU 사용 시 속도를 위해 fp16 사용
)

# 6. Trainer 객체 생성 및 학습 시작
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train_dataset,
    eval_dataset=tokenized_eval_dataset,
    compute_metrics=compute_metrics,
)

print("\n--- 파인 튜닝 시작 ---")
trainer.train()

# 7. 최종 모델 저장
# 학습된 모델과 토크나이저를 지정된 경로에 저장합니다.
print(f"\n--- 파인 튜닝 완료, 모델을 {OUTPUT_DIR}에 저장 중 ---")
trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR) 

print("모델 저장 완료. 이제 저장된 모델을 로드하여 바로 사용할 수 있습니다.")