utils.py

from typing import Optional
import random
from contextlib import contextmanager
import numpy as np
import pandas as pd
import ast
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, PeftModel, get_peft_model, prepare_model_for_kbit_training
import re


device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')

def initialize_seeds(seed_num: int):
    random.seed(seed_num)
    np.random.seed(seed_num)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    torch.manual_seed(seed_num)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed_num)
        torch.cuda.manual_seed_all(seed_num)

def str_to_dict_eedi_df(df: pd.DataFrame):
    cols = ["correct_option", "gt_distractors", "generated_distractors", "log_probs", "distractors", "rank", "proportion"]
    cols = [col for col in cols if col in df.columns]
    for i, row in df.iterrows():
        for col in cols:
            try:
                df.at[i, col] = ast.literal_eval(row[col])
            except Exception:
                df.at[i, col] = None
    return df

bnb_config = BitsAndBytesConfig(
    load_in_8bit=True,
)

peft_config = LoraConfig(
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    r=32,
    lora_alpha=16,
    lora_dropout=0.05,
    task_type="CAUSAL_LM",
    inference_mode=False,
)

def get_base_model(base_model_name: str, tokenizer: AutoTokenizer):
    base_model = AutoModelForCausalLM.from_pretrained(
        base_model_name,
        pad_token_id=tokenizer.pad_token_id,
        quantization_config=bnb_config,
        torch_dtype=torch.float32,
        device_map={"": 0}
    )
    base_model.config.use_cache = False
    base_model.config.pretraining_tp = 1
    return base_model

def get_model(base_model_name: str, model_name: Optional[str], pt_model_name: Optional[str],
              test: bool, use_gradient_checkpointing: bool = True):
    tokenizer = AutoTokenizer.from_pretrained(base_model_name)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"
    model = get_base_model(base_model_name, tokenizer)
    if test and model_name:
        # Note we are loading adapter on quantized model and not merging
        # Recommended here - https://huggingface.co/docs/trl/main/en/dpo_trainer#downsides-to-merging-qlora-before-dpo-approach-2
        # Also prevents empty responses generated by Llama models
        model = PeftModel.from_pretrained(model, model_name)
    else:
        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=use_gradient_checkpointing)
        if pt_model_name:
            model = PeftModel.from_pretrained(model, pt_model_name, is_trainable=True, adapter_name="default")
            # Load same lora weights to separate adapter to be used as reference model
            model.load_adapter(pt_model_name, adapter_name="lora_ref")
        else:
            model = get_peft_model(model, peft_config)
    return model, tokenizer

@contextmanager
def ref_model_ctm(model: PeftModel):
    model.set_adapter("lora_ref")
    model.eval()
    yield
    model.set_adapter("default")
    model.train()

def clean_string(string):
    string = string.lower()
    string = string.replace("$", "")
    string = string.replace("£", "")
    string = string.replace("\\%", "")
    string = string.replace("%", "")
    string = string.replace("...", "\\ldots")
    string = string.replace('÷', '\\div')
    string = string.replace('≥', '\\geq')
    string = string.replace('≤', '\\leq')
    string = string.replace('≠', '\\neq')
    string = string.replace('≈', '\\approx')
    string = string.replace('δ', '\\delta')
    string = string.replace('|', '\\vert')
    string = string.replace(" hours", "")
    string = string.replace("\\(\\frac", "\\frac")
    string = string.replace("}\\)", "}")
    string = string.replace("\\[", "")
    string = string.replace("\\]", "")
    if string == "both towers are the same height":
        string = "$a \\& b$ are the same height"
    if string == 'square root of 96':
        string = '\\sqrt{96}'
    if string == "\\frac{2.0}{3.0}":
        string = "\\frac{2}{3}"

    # convert / and \div fractions to \frac
    string = re.sub(r"([\d\.]+)\s*(/|\\div)\s*([\d\.]+)", r"\\frac{\g<1>}{\g<3>}", string) 
    # convert x to \times
    string = re.sub(r'\s*×\s*', r' \\times ', string)
    # convert √ to \\sqrt{}
    string = re.sub(r'√', r'\\sqrt', string) 
    # convert 2 cm to 2 \mathrm{~cm}
    string = re.sub(r'(\d+(?:\.\d+)?)\s*cm',  r'\1 \\mathrm{~cm}', string)
    # convert 2 m to 2 \mathrm{~m}
    string = re.sub(r'(\d+(?:\.\d+)?)\s*m',  r'\1 \\mathrm{~m}', string)
    # convert 2 km to 2 mathrm{~km}
    string = re.sub(r'(\d+(?:\.\d+)?)\s*km',  r'\1 \\mathrm{~km}', string)

    # convert p^2 to p^{2}
    string = re.sub(r'([a-zA-Z])\^(\d+)', r'\1^{\2}', string)
    
    # remove hyphen between words
    string = re.sub(r'([a-zA-Z]+)-([a-zA-Z]+)', r'\1\2', string)
    
    string = string.replace('\\mathrm{~m}athrm{~cm}', '\\mathrm{~cm}')
    string = string.replace('\\mathrm{~m}ore', 'more')
    string = string.replace(' ', '')
    string = string.strip()
    return string

def relaxed_metric(gt_distractors, generated_distractors):
    # relaxed metric, if any gt_distractors are in the generated distractors, then it is correct
    correct = 0
    for idx, gt_distractor in enumerate(gt_distractors):
        for generated_distractor in generated_distractors[idx]:
            if generated_distractor in gt_distractor:
                correct += 1
                break
    return correct/len(gt_distractors)

def hard_metric(gt_distractors, generated_distractors):
    # hard metric, if all gt_distractors are in the generated distractors, then it is correct
    correct = 0
    for idx, gt_distractor in enumerate(gt_distractors):
        if gt_distractor[0] in generated_distractors[idx] and gt_distractor[1] in generated_distractors[idx] and gt_distractor[2] in generated_distractors[idx]:
            correct += 1
    return correct/len(gt_distractors)

def weighted_proportional_metric(gt_distractors, generated_distractors, proportions):
    correct = 0
    zero_proportion = 0
    for idx, generated_distractor in enumerate(generated_distractors):
        props = 0
        for gt_distractor in gt_distractors[idx]:
            if gt_distractor in generated_distractor:
                props += proportions[idx][gt_distractor]
        sum = proportions[idx][gt_distractors[idx][0]] + proportions[idx][gt_distractors[idx][1]] + proportions[idx][gt_distractors[idx][2]]
        if sum == 0:
            zero_proportion += 1
        else:
            correct += props/sum
    return correct/(len(gt_distractors) - zero_proportion)

def proportional_metric(gt_distractors, generated_distractors):
    correct = 0
    for idx, generated_distractor in enumerate(generated_distractors):
        props = 0
        for gt_distractor in gt_distractors[idx]:
            if gt_distractor in generated_distractor:
                props += 1
        correct += props/3
    return correct/len(gt_distractors)

def process(distractors, question_num, name):
    questions = []
    correct_options = []
    gt_distractors = []
    correct_answers = []
    data = pd.read_csv("data/eedi_test_20_cleaned_2_24_26.csv")
    data = str_to_dict_eedi_df(data)
    for idx, row in data.iterrows():
        questions.append(row["question"].strip())
        correct_options.append(row["correct_option"])
        gt_distractors.append(row["distractors"])
        correct_answers.append([row["correct_option"]['option'].strip()])

    counter = 0
    qualified_questions = []
    qualified_correct_options = []
    qualified_gt_distractors = []
    qualified_generated_distractors = []

    unqualified_questions = []
    unqualified_correct_options = []
    unqualified_gt_distractors = []
    unqualified_generated_distractors = []

    for idx, val in enumerate(distractors):
        # clean the correct answer and distractors
        cleaned_correct_answer = list(map(clean_string, correct_answers[idx]))
        cleaned_distractors = list(map(clean_string, val))
        # remove the correct answer from the distractors based on the cleaned version, but remove the correct answer from distractors[idx]
        # this is the correct way
        ori_distractors = [distractors[idx][i] for i, distractor in enumerate(cleaned_distractors) if distractor not in cleaned_correct_answer]
        # ori_distractors = [distractor for i, distractor in enumerate(cleaned_distractors) if distractor not in cleaned_correct_answer]

        if len(ori_distractors) >= question_num:
            qualified_questions.append(questions[idx])
            qualified_correct_options.append(correct_options[idx])
            qualified_gt_distractors.append(gt_distractors[idx])
            # convert ori_distractors to a dictionary with key is "option"
            qualified_generated_distractors.append([{"option": distractor} for distractor in ori_distractors[:question_num]])
        else: 
            counter += 1
            unqualified_questions.append(questions[idx])
            unqualified_correct_options.append(correct_options[idx])
            unqualified_gt_distractors.append(gt_distractors[idx])
            unqualified_generated_distractors.append([{"option": distractor} for distractor in ori_distractors])
            
    print(f"number of questions with less than {question_num} distractors:", counter)

    qualified_dict = {"question": qualified_questions, "correct_option": qualified_correct_options, "gt_distractors": qualified_gt_distractors, "generated_distractors": qualified_generated_distractors}
    unqualified_dict = {"question": unqualified_questions, "correct_option": unqualified_correct_options, "gt_distractors": unqualified_gt_distractors, "generated_distractors": unqualified_generated_distractors}
        
    qualified_df = pd.DataFrame(qualified_dict)
    unqualified_df = pd.DataFrame(unqualified_dict)
    
    return qualified_df, unqualified_df

def complement_process(distractors, question_num, name):
    questions = []
    correct_options = []
    existing_distractors = []
    gt_distractors = []
    correct_answers = []
    data = pd.read_csv(f"data/{name}_less_than_{question_num}.csv")
    data = str_to_dict_eedi_df(data)
    for idx, row in data.iterrows():
        questions.append(row["question"].strip())
        correct_options.append(row["correct_option"])
        # convert the dictionary into a list, the key is option
        existing_distractors.append([distractor['option'] for distractor in row["generated_distractors"]])
        gt_distractors.append(row["gt_distractors"])
        correct_answers.append([row["correct_option"]['option'].strip()])
    
    generated_distractors = []
    for idx, val in enumerate(distractors):
        # clean the correct answer and distractors
        cleaned_correct_answer = list(map(clean_string, correct_answers[idx]))
        cleaned_distractors = list(map(clean_string, val))
        # remove the correct answer from the distractors based on the cleaned version, but remove the correct answer from distractors[idx]
        ori_distractors = [distractors[idx][i] for i, distractor in enumerate(cleaned_distractors) if distractor not in cleaned_correct_answer]
        # ori_distractors = [distractor for i, distractor in enumerate(cleaned_distractors) if distractor not in cleaned_correct_answer]
        # add original distractors to the existing distractors if original distractors are not in the existing distractors
        for distractor in ori_distractors:
            if distractor not in existing_distractors[idx]:
                existing_distractors[idx].append(distractor)

        if len(existing_distractors[idx]) < question_num:
            print(idx)
            print(questions[idx])
            print(correct_answers[idx])
            print(existing_distractors[idx])
            print(len(existing_distractors[idx]))
            print('_______'*10)
            for i in range(question_num - len(existing_distractors[idx])):
                existing_distractors[idx].append(f"placeholder{i+1}")     
            generated_distractors.append([{"option": distractor} for distractor in existing_distractors[idx]])
        else:
            generated_distractors.append([{"option": distractor} for distractor in existing_distractors[idx][:question_num]])

    dict = {"question": questions, "correct_option": correct_options, "gt_distractors": gt_distractors, "generated_distractors": generated_distractors}
    df = pd.DataFrame(dict)
    return df