from transformers import BertTokenizer, BertForMaskedLM
import torch

# 모델과 토크나이저 로드
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertForMaskedLM.from_pretrained("bert-base-uncased")
model.eval()

# 마스킹된 문장 입력
text = "The capital of France is [MASK]."
inputs = tokenizer(text, return_tensors="pt")

# 예측 수행
with torch.no_grad():
    outputs = model(**inputs)
    predictions = outputs.logits

# [MASK] 위치의 상위 5개 예측 토큰
mask_index = (inputs["input_ids"] == tokenizer.mask_token_id).nonzero(as_tuple=True)[1]
top_5 = torch.topk(predictions[0, mask_index, :], 5, dim=-1)

print("입력:", text)
print("상위 5개 예측:")
for i, (token_id, score) in enumerate(zip(top_5.indices[0], top_5.values[0])):
    token = tokenizer.decode(token_id)
    print(f"  {i+1}. {token} (점수: {score:.4f})")

from transformers import BertTokenizer, BertForSequenceClassification
from transformers import Trainer, TrainingArguments
from datasets import load_dataset

# 데이터셋 로드 (IMDB 감성 분류)
dataset = load_dataset("imdb")
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

# 토큰화 함수
def tokenize_fn(examples):
    """텍스트를 BERT 입력 형식으로 토큰화합니다."""
    return tokenizer(
        examples["text"],
        padding="max_length",
        truncation=True,
        max_length=256,
    )

tokenized = dataset.map(tokenize_fn, batched=True)

# 모델 초기화 (이진 분류)
model = BertForSequenceClassification.from_pretrained(
    "bert-base-uncased",
    num_labels=2,
)

# 학습 설정
training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=32,
    learning_rate=2e-5,
    weight_decay=0.01,
    warmup_ratio=0.1,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
)

# Trainer로 학습 실행
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized["train"],
    eval_dataset=tokenized["test"],
)

trainer.train()

from transformers import BertTokenizer, BertModel
import torch

tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertModel.from_pretrained("bert-base-uncased")
model.eval()

# 문장 인코딩
text = "BERT produces contextualized embeddings."
inputs = tokenizer(text, return_tensors="pt")

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

# [CLS] 토큰의 Hidden State (문장 표현)
cls_embedding = outputs.last_hidden_state[:, 0, :]  # shape: (1, 768)

# 전체 토큰의 Hidden State
all_embeddings = outputs.last_hidden_state  # shape: (1, seq_len, 768)

# 평균 풀링으로 문장 임베딩 생성
# attention_mask를 사용하여 패딩 토큰 제외
mask = inputs["attention_mask"].unsqueeze(-1).float()
mean_embedding = (outputs.last_hidden_state * mask).sum(dim=1) / mask.sum(dim=1)

print(f"[CLS] 임베딩 차원: {cls_embedding.shape}")
print(f"평균 풀링 임베딩 차원: {mean_embedding.shape}")