Views
No views yet
1import numpy as np
2import torch
3from transformers import AutoModelForMultipleChoice, AutoTokenizer
4
5model_name = "JesseStover/L2AI-dictionary-klue-bert-base"
6tokenizer = AutoTokenizer.from_pretrained(model_name)
7model = AutoModelForMultipleChoice.from_pretrained(model_name)
8model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
9
10prompts = "\"강아지는 뽀송뽀송하다.\"에 있는 \"강아지\"의 정의는 "
11candidates = [
12 "\"(명사) 개의 새끼\"예요.",
13 "\"(명사) 부모나 할아버지, 할머니가 자식이나 손주를 귀여워하면서 부르는 말\"이예요."
14]
15
16inputs = tokenizer(
17 [[prompt, candidate] for candidate in candidates],
18 return_tensors="pt",
19 padding=True
20)
21
22labels = torch.tensor(0).unsqueeze(0)
23
24with torch.no_grad():
25 outputs = model(
26 **{k: v.unsqueeze(0) for k, v in inputs.items()}, labels=labels
27 )
28
29print({i: float(x) for i, x in enumerate(outputs.logits.softmax(1)[0])})