XGLM-1.7B

import torch
import torch.nn.functional as F
from openmind import is_torch_npu_available, AutoTokenizer
from transformers import XGLMForCausalLM
 

if is_torch_npu_available():
    device = "npu:0"
elif torch.cuda.is_available():
    device = "cuda:0"
else:
    device = "cpu"

model_name_or_path = 'PyTorch-NPU/xglm_1.7b'
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
model = XGLMForCausalLM.from_pretrained(model_name_or_path, trust_remote_code=True, device_map=device)

data_samples = {
    'en': [
        {
            "premise": "I wanted to conserve energy.",
            "choice1": "I swept the floor in the unoccupied room.",
            "choice2": "I shut off the light in the unoccupied room.",
            "question": "effect",
            "label": "1"
        },
        {
            "premise": "The flame on the candle went out.",
            "choice1": "I blew on the wick.",
            "choice2": "I put a match to the wick.",
            "question": "cause",
            "label": "0"
        }
    ],
    'zh': [
        {
            "premise": "我想节约能源。",
            "choice1": "我在空着的房间里扫了地板。",
            "choice2": "我把空房间里的灯关了。",
            "question": "effect",
            "label": "1"
        },
        {
            "premise": "蜡烛上的火焰熄灭了。",
            "choice1": "我吹灭了灯芯。",
            "choice2": "我把一根火柴放在灯芯上。",
            "question": "cause",
            "label": "0"
        }
    ],
    'hi': [
        {
            "premise": "M te vle konsève enèji.",
            "choice1": "Mwen te fin baleye chanm lib la.",
            "choice2": "Mwen te femen limyè nan chanm lib la.",
            "question": "effect",
            "label": "1"
        },
        {
            "premise": "Flam bouji a te etenn.",
            "choice1": "Mwen te soufle bouji a.",
            "choice2": "Mwen te limen mèch bouji a.",
            "question": "cause",
            "label": "0"
        }
    ]
}


def get_logprobs(prompt, device):
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    input_ids, output_ids = inputs["input_ids"], inputs["input_ids"][:, 1:]
    outputs = model(**inputs, labels=input_ids)
    logits = outputs.logits
    logprobs = torch.gather(F.log_softmax(logits, dim=2), 2,
                            output_ids.unsqueeze(2))
    return logprobs


def COPA_eval(prompt, alternative1, alternative2, device):
    lprob1 = get_logprobs(prompt + "\n" + alternative1, device).sum()
    lprob2 = get_logprobs(prompt + "\n" + alternative2, device).sum()
    return 0 if lprob1 > lprob2 else 1


for lang in data_samples:
    for idx, example in enumerate(data_samples[lang]):
        predict = COPA_eval(example["premise"], example["choice1"],
                            example["choice2"], device)
        print(f'{lang}-{idx}', predict, example['label'])