from datasets import load_dataset

# NSMC 데이터셋 로드
dataset = load_dataset("nsmc")

print(f"학습 데이터: {len(dataset['train']):,}건")  # 150,000건
print(f"테스트 데이터: {len(dataset['test']):,}건")  # 50,000건

# 데이터 구조 확인
print(dataset["train"][0])
# {'id': '9976970', 'document': '아 더빙.. 진짜 짜증나네요 목소리', 'label': 0}

# 라벨 분포 확인
import collections
labels = dataset["train"]["label"]
counter = collections.Counter(labels)
print(f"긍정(1): {counter[1]:,}건, 부정(0): {counter[0]:,}건")
# 긍정(1): 75,173건, 부정(0): 74,827건 → 균형 잡힌 데이터셋

import re
from transformers import AutoTokenizer

# 한국어 PLM 토크나이저 로드
model_name = "monologg/koelectra-base-v3-discriminator"
tokenizer = AutoTokenizer.from_pretrained(model_name)

def clean_text(text: str) -> str:
    """텍스트를 정제합니다."""
    if not isinstance(text, str):
        return ""
    # HTML 태그 제거
    text = re.sub(r"<[^>]+>", "", text)
    # 반복 문자 정규화 (ㅋㅋㅋㅋ → ㅋㅋ)
    text = re.sub(r"(.)\1{3,}", r"\1\1", text)
    # 특수문자 최소화 (마침표, 물음표, 느낌표만 유지)
    text = re.sub(r"[^\w\s.?!]", " ", text)
    # 다중 공백 제거
    text = re.sub(r"\s+", " ", text).strip()
    return text

def preprocess_dataset(examples):
    """데이터셋 전처리 함수 (batched)"""
    # 텍스트 정제
    cleaned = [clean_text(doc) for doc in examples["document"]]

    # 토큰화
    tokenized = tokenizer(
        cleaned,
        padding="max_length",
        truncation=True,
        max_length=128,
        return_tensors="pt",
    )

    tokenized["labels"] = examples["label"]
    return tokenized

# 빈 문서 필터링
dataset = dataset.filter(lambda x: x["document"] is not None and len(str(x["document"])) > 0)

# 전처리 적용
tokenized_dataset = dataset.map(
    preprocess_dataset,
    batched=True,
    remove_columns=dataset["train"].column_names,
)

# PyTorch 텐서 형식 설정
tokenized_dataset.set_format("torch")

print(f"전처리 완료 - 학습: {len(tokenized_dataset['train'])}건")
print(f"입력 형태: {tokenized_dataset['train'][0]['input_ids'].shape}")

# 원본 vs 정제 결과 비교
sample = dataset["train"][100]
print(f"원본: {sample['document']}")
print(f"정제: {clean_text(sample['document'])}")

# 토큰화 결과 확인
tokens = tokenizer.tokenize(clean_text(sample["document"]))
print(f"토큰: {tokens[:20]}...")
print(f"토큰 수: {len(tokens)}")

import torch
from transformers import (
    AutoModelForSequenceClassification,
    TrainingArguments,
    Trainer,
)
import evaluate
import numpy as np

# 모델 로드 (분류 헤드 포함)
model = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    num_labels=2,  # 긍정/부정 2클래스
)

# 디바이스 확인
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"학습 디바이스: {device}")

# 평가 지표 설정
accuracy_metric = evaluate.load("accuracy")
f1_metric = evaluate.load("f1")

def compute_metrics(eval_pred):
    """학습 중 평가 지표를 계산합니다."""
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)

    accuracy = accuracy_metric.compute(
        predictions=predictions, references=labels
    )
    f1 = f1_metric.compute(
        predictions=predictions, references=labels, average="binary"
    )

    return {
        "accuracy": accuracy["accuracy"],
        "f1": f1["f1"],
    }

# 학습 설정
training_args = TrainingArguments(
    output_dir="./results/koelectra-nsmc",
    num_train_epochs=3,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=64,
    learning_rate=5e-5,
    weight_decay=0.01,
    warmup_ratio=0.1,
    eval_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    logging_steps=100,
    fp16=torch.cuda.is_available(),  # GPU 사용 시 Mixed Precision
    report_to="none",  # wandb 비활성화 (필요 시 "wandb"로 변경)
)

# Trainer 생성
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],
    compute_metrics=compute_metrics,
)

# 학습 시작
print("Fine-tuning 시작...")
train_result = trainer.train()

# 학습 결과
print(f"\n학습 완료:")
print(f"  총 학습 시간: {train_result.metrics['train_runtime']:.0f}초")
print(f"  최종 손실: {train_result.metrics['train_loss']:.4f}")

import numpy as np
from sklearn.metrics import (
    classification_report,
    confusion_matrix,
    ConfusionMatrixDisplay,
)
import matplotlib.pyplot as plt

# 테스트 세트 평가
eval_results = trainer.evaluate()
print(f"테스트 Accuracy: {eval_results['eval_accuracy']:.4f}")
print(f"테스트 F1 Score: {eval_results['eval_f1']:.4f}")

# 예측 수행
predictions = trainer.predict(tokenized_dataset["test"])
pred_labels = np.argmax(predictions.predictions, axis=-1)
true_labels = predictions.label_ids

# 상세 분류 리포트
print("\n분류 리포트:")
print(classification_report(
    true_labels,
    pred_labels,
    target_names=["부정", "긍정"],
))

# Confusion Matrix 시각화
cm = confusion_matrix(true_labels, pred_labels)
disp = ConfusionMatrixDisplay(
    confusion_matrix=cm,
    display_labels=["부정", "긍정"],
)
fig, ax = plt.subplots(figsize=(6, 5))
disp.plot(ax=ax, cmap="Blues")
ax.set_title("감성 분석 Confusion Matrix")
plt.tight_layout()
plt.savefig("confusion_matrix.png", dpi=150)
plt.show()

# 오분류 사례 분석
misclassified = []
original_test = dataset["test"]

for i, (pred, true) in enumerate(zip(pred_labels, true_labels)):
    if pred != true:
        misclassified.append({
            "text": original_test[i]["document"],
            "true_label": "긍정" if true == 1 else "부정",
            "pred_label": "긍정" if pred == 1 else "부정",
        })

print(f"\n오분류 수: {len(misclassified)}건 / {len(true_labels)}건")
print("\n오분류 예시 (상위 5건):")
for item in misclassified[:5]:
    print(f"  [{item['true_label']}→{item['pred_label']}] {item['text'][:60]}...")

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# 모델 저장
model_path = "./saved_model/koelectra-nsmc"
trainer.save_model(model_path)
tokenizer.save_pretrained(model_path)
print(f"모델 저장 완료: {model_path}")

# 추론 함수
def predict_sentiment(text: str, model_path: str = model_path) -> dict:
    """텍스트의 감성을 예측합니다."""
    # 모델 로드
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSequenceClassification.from_pretrained(model_path)
    model.eval()

    # 전처리 및 토큰화
    cleaned = clean_text(text)
    inputs = tokenizer(
        cleaned,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=128,
    )

    # 추론
    with torch.no_grad():
        outputs = model(**inputs)
        probabilities = torch.softmax(outputs.logits, dim=-1)

    neg_prob = probabilities[0][0].item()
    pos_prob = probabilities[0][1].item()
    predicted_label = "긍정" if pos_prob > neg_prob else "부정"

    return {
        "text": text,
        "sentiment": predicted_label,
        "confidence": max(pos_prob, neg_prob),
        "probabilities": {
            "긍정": round(pos_prob, 4),
            "부정": round(neg_prob, 4),
        }
    }

# 추론 테스트
test_texts = [
    "이 영화 진짜 재밌었어요! 다음 편도 기대됩니다.",
    "시간 낭비였습니다. 스토리가 너무 허술해요.",
    "보통이었어요. 나쁘지는 않지만 특별하지도 않았습니다.",
]

for text in test_texts:
    result = predict_sentiment(text)
    print(f"[{result['sentiment']}] ({result['confidence']:.2%}) {text}")

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# 모델 로드 (전역으로 한 번만)
model_path = "./saved_model/koelectra-nsmc"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)
model.eval()

def analyze_sentiment(text: str) -> dict:
    """Gradio용 감성 분석 함수"""
    if not text.strip():
        return {"긍정": 0.0, "부정": 0.0}

    cleaned = clean_text(text)
    inputs = tokenizer(
        cleaned,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=128,
    )

    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.softmax(outputs.logits, dim=-1)

    return {
        "긍정": round(probs[0][1].item(), 4),
        "부정": round(probs[0][0].item(), 4),
    }

# Gradio 인터페이스
demo = gr.Interface(
    fn=analyze_sentiment,
    inputs=gr.Textbox(
        label="리뷰 텍스트",
        placeholder="분석할 한국어 리뷰를 입력하세요...",
        lines=3,
    ),
    outputs=gr.Label(
        label="감성 분석 결과",
        num_top_classes=2,
    ),
    title="한국어 감성 분석기",
    description=(
        "KoELECTRA 기반 한국어 영화 리뷰 감성 분석 모델입니다. "
        "NSMC 데이터셋으로 Fine-tuning되었습니다."
    ),
    examples=[
        ["정말 감동적인 영화였습니다. 눈물이 멈추지 않았어요."],
        ["연출도 엉성하고 배우 연기도 최악이었습니다."],
        ["그럭저럭 볼만했어요. 시간 때우기 좋습니다."],
    ],
    theme=gr.themes.Soft(),
)

# 실행
demo.launch(share=False)  # share=True로 공개 URL 생성 가능