import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# 1. 모델이 저장된 폴더 경로 지정
LOAD_MODEL_PATH = "./xtreme-distil-review-classifier"

# 2. GPU/CPU 장치 설정
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"사용 장치: {device}")

# 3. 저장된 토크나이저와 모델 로드
# 저장된 config.json과 model.safetensors 파일을 바탕으로 로드합니다.
print(f"\n--- 모델 로드 중: {LOAD_MODEL_PATH} ---")
loaded_tokenizer = AutoTokenizer.from_pretrained(LOAD_MODEL_PATH)
loaded_model = AutoModelForSequenceClassification.from_pretrained(LOAD_MODEL_PATH)

# 모델을 설정된 장치(GPU 또는 CPU)로 이동
loaded_model.to(device)
loaded_model.eval() # 모델을 평가 모드로 설정 (필수)

# 4. 분류 함수 정의
def classify_review(text):
    # 텍스트를 토큰화하고 장치로 이동
    inputs = loaded_tokenizer(
        text,
        return_tensors="pt",  # PyTorch 텐서로 반환
        padding=True,
        truncation=True
    ).to(device)
    
    # 모델 추론 (Inference)
    with torch.no_grad():
        outputs = loaded_model(**inputs)
        
    # 결과 처리
    probabilities = torch.softmax(outputs.logits, dim=1)
    predicted_class_id = probabilities.argmax().item()
    
    # 레이블 매핑 (파인 튜닝 시 설정한 0: 부정, 1: 긍정 기준)
    label_map = {0: "부정 (Negative)", 1: "긍정 (Positive)"}
    predicted_label = label_map[predicted_class_id]
    confidence = probabilities[0][predicted_class_id].item()
    
    return predicted_label, confidence

# 5. 새로운 당근마켓 리뷰 테스트 실행
new_reviews = [
    "매너가 정말 좋으세요! 기분 좋은 거래네요",
    "물건 상태가 생각보다 너무 안 좋아서 속았다는 느낌이 듭니다.",
    "빠른 거래 감사합니다. 문제 없이 잘 받았어요.",
    "연락을 안받네요",
]

print("\n--- 새로운 리뷰 분류 결과 ---")
for review in new_reviews:
    label, confidence = classify_review(review)
    print(f"리뷰: '{review}'")
    print(f"  -> 예측 분류: **{label}** (확률: {confidence:.4f})")
    print("-" * 35)