likelihood_eval.py

"""nusacrowd zero-shot prompt.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1Ru8DyS2ALWfRdkjOPHj-KNjw6Pfa44Nd
"""
import os, sys
import csv
import argparse
from glob import glob

from numpy import argmax, stack
import pandas as pd
from tqdm import tqdm
from sklearn.metrics import classification_report, precision_recall_fscore_support

import torch
import torch.nn.functional as F
import datasets

from peft import PeftModel
from transformers import AutoConfig, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM, set_seed, BitsAndBytesConfig

from utils.prompts import CONFIG_TO_PROMPT

DEBUG=False

def to_prompt(prompt, ff_word, l1_sent, l2_sent, sentence_query, eng_meaning, meaning_l1, meaning_l2):        
    return prompt.replace('[FF]', ff_word).replace('[L1]', l1_sent).replace('[L2]', l2_sent).replace('[L]', sentence_query).replace('[English]', eng_meaning).replace('[Meaning_L1]', meaning_l1).replace('[Meaning_L2]', meaning_l2)

@torch.inference_mode()
def get_logprobs(model, tokenizer, inputs, label_ids=None, label_attn=None):
    inputs = tokenizer(inputs, return_tensors="pt", padding=True, truncation=True, max_length=1024).to('cuda')    
    if model.config.is_encoder_decoder:
        label_ids, label_attn = label_ids.to('cuda'), label_attn.to('cuda')
        logits = model(**inputs, labels=label_ids, decoder_attention_mask=label_attn).logits
        logprobs = torch.gather(F.log_softmax(logits, dim=-1), 2, label_ids.unsqueeze(2)).squeeze(dim=-1)
        return logprobs.sum(dim=-1).cpu()
    else:
        logits = model(**inputs).logits
        output_ids = inputs["input_ids"][:, 1:]
        logprobs = torch.gather(F.log_softmax(logits, dim=-1), 2, output_ids.unsqueeze(2)).squeeze(dim=-1)
        logprobs[inputs["attention_mask"][:, :-1] == 0] = 0
        return logprobs.sum(dim=1).cpu()

@torch.inference_mode()
def predict_classification(model, tokenizer, prompts, label_names):
    if type(label_names) is tuple:
        if model.config.is_encoder_decoder:            
            batched_prompts, batched_labels = [], []
            for prompt, label_group in zip(prompts, list(zip(*label_names))):
                for label in label_group:
                    batched_prompts.append(prompt.replace('[LABELS_CHOICE]', ''))
                    batched_labels.append(label)
            labels_encoded = tokenizer(batched_labels, add_special_tokens=False, padding=True, return_tensors='pt')
            batched_label_ids, batched_label_attn = labels_encoded['input_ids'], labels_encoded['attention_mask']
            probs = get_logprobs(model, tokenizer, batched_prompts, batched_label_ids, batched_label_attn).view(len(prompts), -1).numpy()
        else:
            batched_prompts = []
            for prompt, label_group in zip(prompts, list(zip(*label_names))):
                for label in label_group:
                    batched_prompts.append(prompt.replace('[LABELS_CHOICE]', label))
            probs = get_logprobs(model, tokenizer, batched_prompts).view(len(prompts), -1).numpy()
    else:
        if model.config.is_encoder_decoder:
            labels_encoded = tokenizer(label_names, add_special_tokens=False, padding=True, return_tensors='pt')
            list_label_ids, list_label_attn = labels_encoded['input_ids'], labels_encoded['attention_mask']

            batched_prompts, batched_label_ids, batched_label_attn = [], [], []
            for prompt in prompts:
                for (label_ids, label_attn) in zip(list_label_ids, list_label_attn):
                    batched_prompts.append(prompt.replace('[LABELS_CHOICE]', ''))
                    batched_label_ids.append(label_ids)
                    batched_label_attn.append(label_attn)
            batched_label_ids = torch.stack(batched_label_ids, dim=0)
            batched_label_attn = torch.stack(batched_label_attn, dim=0)
            probs = get_logprobs(model, tokenizer, batched_prompts, batched_label_ids, batched_label_attn).view(len(prompts), -1).numpy()
        else:
            batched_prompts = []
            for prompt in prompts:
                for label in label_names:
                    batched_prompts.append(prompt.replace('[LABELS_CHOICE]', label))
            probs = get_logprobs(model, tokenizer, batched_prompts).view(len(prompts), -1).numpy()
    return probs

if __name__ == '__main__':
    ##
    # Set seed
    ##
    set_seed(42)

    ##
    # Extract the arguments from the argparse
    ##
    parser = argparse.ArgumentParser(description='Running likelihood-based classification for cognate benchmark')
    parser.add_argument('model_path_or_name', help='Path to the model directory or name of the pre-trained model')
    parser.add_argument('batch_size', type=int, default=128, help='Batch size for model inference')
    parser.add_argument('save_every', type=int, default=16, help='Save every k iterations')
    args = parser.parse_args()
    
    model_path_or_name = args.model_path_or_name
    batch_size = args.batch_size
    save_every = args.save_every

    print(f'MODEL_NAME: {model_path_or_name} | BATCH_SIZE: {batch_size} | SAVE_EVERY: {save_every}')
    
    out_dir = './outputs_likelihood'
    metric_dir = './metrics_likelihood'
    os.makedirs(out_dir, exist_ok=True) 
    os.makedirs(metric_dir, exist_ok=True) 

    if os.path.exists(f'{metric_dir}/results_{model_path_or_name.split("/")[-1]}.csv'):
        print(f'Skipping {metric_dir}/results_{model_path_or_name.split("/")[-1]}.csv')
        sys.exit(0)

    ##
    # Load Model
    ##
    config = AutoConfig.from_pretrained(model_path_or_name, trust_remote_code=True)
    
    tokenizer = AutoTokenizer.from_pretrained(model_path_or_name, truncation_side='left', padding_side='left', trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token if tokenizer.eos_token is not None else tokenizer.bos_token
    
    if not config.is_encoder_decoder:
        model = AutoModelForCausalLM.from_pretrained(
            model_path_or_name, device_map="auto", quantization_config=BitsAndBytesConfig(load_in_8bit=True), trust_remote_code=True
        )
    else:
        model = AutoModelForSeq2SeqLM.from_pretrained(
            model_path_or_name, device_map="auto", quantization_config=BitsAndBytesConfig(load_in_8bit=True), trust_remote_code=True
        )
    model.eval()

    ##
    # Run Inference
    ##
    metrics, count = [], 0
    for config_name in ['id_tl', 'id_tl_common', 'zh_ja', 'zh_ja_common', 'id_ms', 'id_ms_common', 'en_de', 'en_de_common']:
        dset = datasets.load_dataset("StingrayBench/StingrayBench", config_name, split=datasets.Split.TEST)
        for task in ["semantic_correctness", "usage_correctness_l1", "usage_correctness_l2"]:
            if 'common' in config_name and 'meaning' in task:
                continue

            print(f'Running inference for {config_name} - {task}')
            prompt_template = CONFIG_TO_PROMPT[task]
            
            # check saved data
            inputs, preds, golds = [], [], []
            if os.path.exists(f'{out_dir}/{config_name}_{task}_{model_path_or_name.split("/")[-1]}.csv'):
                print("Output exist, use partial log instead")
                with open(f'{out_dir}/{config_name}_{task}_{model_path_or_name.split("/")[-1]}.csv') as csvfile:
                    reader = csv.DictReader(csvfile)
                    for row in reader:
                        inputs.append(row["Input"])
                        preds.append(row["Pred"])
                        golds.append(row["Gold"])
                print(f"Skipping until {len(preds)}")
                
            for idx in range(len(preds), len(dset), batch_size):
                lang1_sent = dset[idx:idx+batch_size]['lang1_sentence']
                lang2_sent = dset[idx:idx+batch_size]['lang2_sentence']
                eng_meaning = dset[idx:idx+batch_size]['lang_en_sentence']
                lang_flag = dset[idx:idx+batch_size]['language_flag']
                meaning_l1 = dset[idx:idx+batch_size]['meaning_l1']
                meaning_l2 = dset[idx:idx+batch_size]['meaning_l2']
                ff_word = dset[idx:idx+batch_size]['word']
                if task == "semantic_correctness":
                    labels = list(map(lambda x: chr(ord('A') + int(x)), dset[idx:idx+batch_size]['semantic_appropriate_label']))
                    label_names = ['A', 'B', 'C']
                elif task == 'usage_correctness_l1':
                    labels = dset[idx:idx+batch_size]['usage_correctness_lang1_answer']
                    label_names = ['Yes', 'No']
                elif task == 'usage_correctness_l2':
                    labels = dset[idx:idx+batch_size]['usage_correctness_lang2_answer']
                    label_names = ['Yes', 'No']
                elif task == 'meaning_l1':
                    labels = list(map(lambda x: str(x), dset[idx:idx+batch_size]['meaning_l1']))
                    false_labels = list(map(lambda x: str(x), dset[idx:idx+batch_size]['meaning_l2']))
                    label_names = (labels, false_labels)
                elif task == 'meaning_l2':
                    labels = list(map(lambda x: str(x), dset[idx:idx+batch_size]['meaning_l2']))
                    false_labels = list(map(lambda x: str(x), dset[idx:idx+batch_size]['meaning_l1']))
                    label_names = (labels, false_labels)
                elif task == 'sentence_meaning':
                    labels = list(map(lambda x: 'A' if x == 'L1' else 'B', dset[idx:idx+batch_size]['sentence_meaning_answer']))
                    label_names = ['A', 'B', 'C']
                elif task == 'word_meaning':
                    labels = list(map(lambda x: 'A' if x == 'L1' else 'B', dset[idx:idx+batch_size]['language_flag']))
                    label_names = ['A', 'B', 'C']
                else: 
                    raise ValueError(f'Unknown task `{task}`')

                # sample prompt
                print("= SAMPLE PROMPT =")
                print(to_prompt(
                    prompt_template, str(ff_word[0]), str(lang1_sent[0]), str(lang2_sent[0]), 
                    lang1_sent[0] if lang_flag[0] == 'L1' else lang2_sent[0], str(eng_meaning[0]), str(meaning_l1[0]), str(meaning_l2[0]))
                 )
                print("\n")

                # zero-shot inference 
                prompts = []
                for i in range(len(ff_word)):
                    prompt_text = to_prompt(
                        prompt_template, str(ff_word[i]), str(lang1_sent[i]), str(lang2_sent[i]), 
                        str(lang1_sent[i]) if lang_flag[i] == 'L1' else str(lang2_sent[i]), str(eng_meaning[i]), str(meaning_l1[i]), str(meaning_l2[i])
                    )
                    prompts.append(f'{prompt_text}\nAnswer: [LABELS_CHOICE]')

                # batch Inference
                out = predict_classification(model, tokenizer, prompts, label_names)
                hyps = argmax(out, axis=-1).tolist()
                if type(label_names) is tuple:
                    for (prompt_text, hyp, label, false_label) in zip(prompts, hyps, labels, false_labels):
                        inputs.append(prompt_text)
                        preds.append(label if hyp == 0 else false_label)
                        golds.append(label)
                else:
                    for (prompt_text, hyp, label) in zip(prompts, hyps, labels):
                        inputs.append(prompt_text)
                        preds.append(label_names[hyp])
                        golds.append(label)
                count += 1
                            
                if count == save_every:
                    # partial saving
                    inference_df = pd.DataFrame(list(zip(inputs, preds, golds)), columns =["Input", 'Pred', 'Gold'])
                    inference_df.to_csv(f'{out_dir}/{config_name}_{task}_{model_path_or_name.split("/")[-1]}.csv', index=False)
                    count = 0

            # saving final
            inference_df = pd.DataFrame(list(zip(inputs, preds, golds)), columns =["Input", 'Pred', 'Gold'])
            inference_df.to_csv(f'{out_dir}/{config_name}_{task}_{model_path_or_name.split("/")[-1]}.csv', index=False)

            cls_report = classification_report(golds, preds, output_dict=True)
            micro_f1, micro_prec, micro_rec, _ = precision_recall_fscore_support(golds, preds, average='micro')
            print('config_name', config_name)
            print('task', task)
            print('accuracy', cls_report['accuracy'])
            print('f1 micro', micro_f1)
            print('f1 macro', cls_report['macro avg']['f1-score'])
            print('f1 weighted', cls_report['weighted avg']['f1-score'])
            print("===\n\n")       

            metrics.append({
                'config_name': config_name,
                'task': task,
                'model': model_path_or_name.split("/")[-1],
                'accuracy': cls_report['accuracy'], 
                'micro_prec': micro_prec,
                'micro_rec': micro_rec,
                'micro_f1_score': micro_f1,
                'macro_prec': cls_report['macro avg']['precision'],
                'macro_rec': cls_report['macro avg']['recall'],
                'macro_f1_score': cls_report['macro avg']['f1-score'],
                'weighted_prec': cls_report['weighted avg']['precision'],
                'weighted_rec': cls_report['weighted avg']['recall'],
                'weighted_f1_score': cls_report['weighted avg']['f1-score'],
            })

    pd.DataFrame(metrics).reset_index().to_csv(f'{metric_dir}/results_{model_path_or_name.split("/")[-1]}.csv', index=False)