from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch

# KLUE-BERT로 Masked LM 추론
model_name = "klue/bert-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForMaskedLM.from_pretrained(model_name)
model.eval()

# 마스킹된 문장
text = "오늘 날씨가 정말 [MASK]."
inputs = tokenizer(text, return_tensors="pt")

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

# [MASK] 위치 예측
mask_idx = (inputs["input_ids"] == tokenizer.mask_token_id).nonzero(as_tuple=True)[1]
top_5 = torch.topk(outputs.logits[0, mask_idx, :], 5, dim=-1)

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

from transformers import AutoTokenizer, AutoModelForSequenceClassification
from transformers import Trainer, TrainingArguments
from datasets import Dataset
import torch

# KoELECTRA로 감성 분류
model_name = "monologg/koelectra-base-v3-discriminator"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    num_labels=2,  # 긍정/부정
)

# 예시 데이터셋 구성
train_data = {
    "text": [
        "이 영화 정말 재미있었어요!",
        "배우들의 연기가 훌륭합니다.",
        "시간 낭비였어요. 별로입니다.",
        "스토리가 너무 지루해서 잠이 왔어요.",
    ],
    "label": [1, 1, 0, 0],  # 1: 긍정, 0: 부정
}
dataset = Dataset.from_dict(train_data)

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

tokenized = dataset.map(tokenize_fn, batched=True)
tokenized.set_format("torch", columns=["input_ids", "attention_mask", "label"])

# 학습 설정
training_args = TrainingArguments(
    output_dir="./ko-sentiment",
    num_train_epochs=5,
    per_device_train_batch_size=4,
    learning_rate=3e-5,
    weight_decay=0.01,
    warmup_ratio=0.1,
    logging_steps=10,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized,
)

trainer.train()

from transformers import pipeline

# KLUE-RoBERTa를 사용한 문장 유사도
classifier = pipeline(
    "text-classification",
    model="klue/roberta-base",
    tokenizer="klue/roberta-base",
)

# 추론 테스트
texts = [
    "오늘 서울 날씨가 맑습니다.",
    "인공지능 기술이 빠르게 발전하고 있습니다.",
    "맛있는 커피를 마시며 코딩하고 있습니다.",
]

for text in texts:
    result = classifier(text)
    print(f"입력: {text}")
    print(f"예측: {result}")
    print()

from transformers import AutoTokenizer, AutoModelForTokenClassification
import torch

# KLUE-BERT NER 모델
model_name = "klue/bert-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# NER 태그 정의 (KLUE NER 기준)
ner_tags = ["O", "B-PS", "I-PS", "B-LC", "I-LC", "B-OG", "I-OG",
            "B-DT", "I-DT", "B-TI", "I-TI", "B-QT", "I-QT"]

model = AutoModelForTokenClassification.from_pretrained(
    model_name,
    num_labels=len(ner_tags),
)

# 문장 토큰화 및 추론
text = "삼성전자가 서울에서 AI 컨퍼런스를 개최합니다."
inputs = tokenizer(text, return_tensors="pt")

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

# 예측 태그 변환
predictions = torch.argmax(outputs.logits, dim=-1)[0]
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])

print(f"입력: {text}")
print("토큰별 NER 태그:")
for token, pred in zip(tokens, predictions):
    tag = ner_tags[pred.item()]
    if tag != "O":  # O 태그가 아닌 것만 출력
        print(f"  {token}: {tag}")