import torch
from transformers import AutoTokenizer, AutoModel
import numpy as np

# 1. GPU/CPU 장치 설정
# CUDA (GPU) 사용 가능하면 'cuda', 아니면 'cpu'로 설정합니다.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"사용 장치: {device}")

# 2. 모델 및 토크나이저 로드
# XLM-RoBERTa-base는 Sequence Classification이 아닌, 일반 임베딩 추출 모델로 로드합니다.
MODEL_NAME = "FacebookAI/xlm-roberta-base"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModel.from_pretrained(MODEL_NAME)

# 모델을 설정된 장치(GPU 또는 CPU)로 이동
model.to(device)

# 3. 임베딩(Embedding) 추출 함수 정의
def get_text_embedding(text):
    # 텍스트를 토큰화하고 장치로 이동
    inputs = tokenizer(
        text,
        return_tensors="pt", # PyTorch 텐서로 반환
        padding=True,
        truncation=True
    ).to(device)
    
    # 모델 추론 (Inference)
    with torch.no_grad():
        # output에는 last_hidden_state (각 토큰의 임베딩) 등이 포함됩니다.
        outputs = model(**inputs)
        
    # 문장 임베딩 추출: [CLS] 토큰의 임베딩을 사용합니다.
    # last_hidden_state의 첫 번째 토큰 (인덱스 0)이 [CLS] 토큰이며, 전체 문장을 대표합니다.
    # shape: (1, 768)
    sentence_embedding = outputs.last_hidden_state[:, 0, :].cpu().numpy()
    
    return sentence_embedding[0] # numpy 배열로 반환

# 4. 당근마켓 리뷰 예제 실행
review_sentences = [
    "매너가 정말 좋으시고 물건도 깨끗해서 만족스러웠어요.",
    "이건 좀 아닌듯. 물건 상태도 별로고 답변도 느렸습니다.",
    "이 모델의 중고 시세는 얼마인가요?", # 일반적인 질문 문장
    "This is a great product for the price." # 외국어 문장도 처리 가능
]

print("\n--- XLM-RoBERTa 임베딩 추출 결과 ---")
for sentence in review_sentences:
    embedding = get_text_embedding(sentence)
    
    print(f"문장: '{sentence}'")
    print(f"  -> 임베딩 차원: {embedding.shape}") # 768차원
    print(f"  -> 임베딩 벡터 일부 (첫 5개): {embedding[:5].round(4)}")
    print("-" * 20)