1import torch
2import torch.nn.functional as F
3
4from transformers import XGLMTokenizer, XGLMForCausalLM
5
6tokenizer = XGLMTokenizer.from_pretrained("facebook/xglm-2.9B")
7model = XGLMForCausalLM.from_pretrained("facebook/xglm-2.9B")
8
9data_samples = {
10 'en': [
11 {
12 "premise": "I wanted to conserve energy.",
13 "choice1": "I swept the floor in the unoccupied room.",
14 "choice2": "I shut off the light in the unoccupied room.",
15 "question": "effect",
16 "label": "1"
17 },
18 {
19 "premise": "The flame on the candle went out.",
20 "choice1": "I blew on the wick.",
21 "choice2": "I put a match to the wick.",
22 "question": "cause",
23 "label": "0"
24 }
25 ],
26 'zh': [
27 {
28 "premise": "我想节约能源。",
29 "choice1": "我在空着的房间里扫了地板。",
30 "choice2": "我把空房间里的灯关了。",
31 "question": "effect",
32 "label": "1"
33 },
34 {
35 "premise": "蜡烛上的火焰熄灭了。",
36 "choice1": "我吹灭了灯芯。",
37 "choice2": "我把一根火柴放在灯芯上。",
38 "question": "cause",
39 "label": "0"
40 }
41 ],
42 'hi': [
43 {
44 "premise": "M te vle konsève enèji.",
45 "choice1": "Mwen te fin baleye chanm lib la.",
46 "choice2": "Mwen te femen limyè nan chanm lib la.",
47 "question": "effect",
48 "label": "1"
49 },
50 {
51 "premise": "Flam bouji a te etenn.",
52 "choice1": "Mwen te soufle bouji a.",
53 "choice2": "Mwen te limen mèch bouji a.",
54 "question": "cause",
55 "label": "0"
56 }
57 ]
58}
59
60def get_logprobs(prompt):
61 inputs = tokenizer(prompt, return_tensors="pt")
62 input_ids, output_ids = inputs["input_ids"], inputs["input_ids"][:, 1:]
63 outputs = model(**inputs, labels=input_ids)
64 logits = outputs.logits
65 logprobs = torch.gather(F.log_softmax(logits, dim=2), 2, output_ids.unsqueeze(2))
66 return logprobs
67
68# Zero-shot evaluation for the Choice of Plausible Alternatives (COPA) task.
69# A return value of 0 indicates that the first alternative is more plausible,
70# while 1 indicates that the second alternative is more plausible.
71def COPA_eval(prompt, alternative1, alternative2):
72 lprob1 = get_logprobs(prompt + "\n" + alternative1).sum()
73 lprob2 = get_logprobs(prompt + "\n" + alternative2).sum()
74 return 0 if lprob1 > lprob2 else 1
75
76for lang in data_samples_long:
77 for idx, example in enumerate(data_samples_long[lang]):
78 predict = COPA_eval(example["premise"], example["choice1"], example["choice2"])
79 print(f'{lang}-{idx}', predict, example['label'])
80
81# en-0 1 1
82# en-1 0 0
83# zh-0 1 1
84# zh-1 0 0
85# hi-0 1 1
86# hi-1 0 0
