run_cont.py

import warnings
from collections import Counter
import itertools
from sentence_transformers import SentenceTransformer, util
from data import IndexingTrainDataset, IndexingCollator
from peft import get_peft_model, LoraConfig, TaskType, PeftModel
from inspect import signature
from transformers import (
    TrainingArguments,
    TrainerCallback,
    MT5Tokenizer,
    MT5TokenizerFast,
    MT5ForConditionalGeneration,
    HfArgumentParser,
    set_seed,
)
from trainer import DSITrainer, MapTrainer
import numpy as np
import torch
import wandb
from dataclasses import dataclass, field
from typing import Optional
import json
# wandb.init(mode='offline')
set_seed(313)

@dataclass
class RunArguments:
    model_name: str = field(default="google/mt5-base")
    base_model_path: str = field(default="../model/mt5-base")
    max_length: Optional[int] = field(default=64)
    id_max_length: Optional[int] = field(default=20)
    train_file: str = field(default="./dataset/fetaqa_inc/train_0.json")
    valid_file: str = field(default="./dataset/fetaqa_inc/test_0.json")
    all_file: str = field(default="./dataset/fetaqa_inc/train_0.json", metadata={"help": "the valid tabids"})
    task: str = field(default="Search",  metadata={"help": "Index, Search"})
    top_k: Optional[int] = field(default=10)
    num_return_sequences: Optional[int] = field(default=10)
    q_max_length: Optional[int] = field(default=32)
    remove_prompt: Optional[bool] = field(default=True)
    peft: Optional[bool] = field(default=False)


@dataclass
class CustomTrainingArguments(TrainingArguments):
    num: int = field(default=2) # num of models, including base and lora
    LoRA_1: str = field(default="./model/feta_LoRA_d1/checkpoint-100")
    LoRA_2: str = field(default="None")
    LoRA_3: str = field(default="None")
    partition_0: str = field(default="./dataset/fetaqa_inc/train_0.json")
    partition_1: str = field(default="./dataset/fetaqa_inc/train_1.json")
    partition_2: str = field(default="None")
    partition_3: str = field(default="None")
    num_train_epochs: int = field(default=3)
    learning_rate: float = field(default=0.0005)
    warmup_steps: float = field(default=10000)
    per_device_train_batch_size: int = field(default=64)
    per_device_eval_batch_size: int = field(default=32)
    evaluation_strategy: str = field(default="steps")
    eval_steps: int = field(default=100)
    max_steps: int = field(default=800)
    save_strategy: str = field(default="steps")
    dataloader_num_workers: int = field(default=6)
    save_steps: int = field(default=100)
    save_total_limit: int = field(default=3)
    gradient_accumulation_steps : int = field(default=6)
    report_to: str = field(default="wandb")
    logging_steps: int = field(default=100)
    metric_for_best_model: str = field(default="Hits@5")
    dataloader_drop_last: bool = field(default=False)
    greater_is_better: bool = field(default=True)
    run_name: str = field(default="feta_inc")
    output_dir: str = field(default="./model/feta_inc0")

class PrintGradientsCallback(TrainerCallback):
    def on_epoch_end(self, args, state, control, **kwargs):
        model = kwargs['model']
        print(f"Epoch {state.epoch} ended.")
        for name, param in model.named_parameters():
            if param.requires_grad:
                print(f"Parameter Name: {name}")
                print("Parameter Value:")
                print(param.data)
                if param.grad is not None:
                    print("Parameter Gradient:")
                    print(param.grad)
                else:
                    print("Parameter Gradient: None")

def find_similar_embeddings(embeddings_1, embeddings_2, topk =1):
    similarity_matrix = np.dot(embeddings_1, embeddings_2.T)
    topk_indices = np.argsort(similarity_matrix, axis=1)[:, -topk:]
    # Sort the top-k indices for each row in ascending order
    topk_indices = np.sort(topk_indices)
    # max_indices = np.argmax(similarity_matrix, axis=1)
    return topk_indices

def inference_on_test_set(model, test_args, test_dataset, tokenizer, restrict_decode_vocab, id_max_length):
    trainer = MapTrainer(
        model=model,
        args=test_args,
        tokenizer=tokenizer,
        data_collator=IndexingCollator(tokenizer=tokenizer),
        restrict_decode_vocab=restrict_decode_vocab,
        id_max_length=id_max_length
    )

    results = trainer.predict(test_dataset=test_dataset)

    predictions = results.predictions
    labels = results.label_ids

    return predictions, labels

def make_compute_metrics(tokenizer, valid_ids):

    def compute_metrics(eval_preds):
        hit_at_1 = 0
        hit_at_5 = 0
        i = 0
        gt2err = dict()
        for beams, label in zip(eval_preds.predictions, eval_preds.label_ids):
            rank_list = tokenizer.batch_decode(beams, skip_special_tokens=True)

            label_id = tokenizer.decode(label, skip_special_tokens=True)
            # filter out duplicates and invalid docids
            filtered_rank_list = []
            for docid in rank_list:
                if docid not in filtered_rank_list and docid in valid_ids:
                    filtered_rank_list.append(docid)
            hits = np.where(np.array(filtered_rank_list)[:5] == label_id)[0]
            # hits = np.where(np.array(filtered_rank_list)[:10] == label_id)[0]

            if len(hits) != 0:
                hit_at_5 += 1
                if hits[0] == 0:
                    hit_at_1 += 1
            i += 1
        return {"Hits@1": hit_at_1 / len(eval_preds.predictions), "Hits@5": hit_at_5 / len(eval_preds.predictions)}
    return compute_metrics

def valid_id(train_pth):
    valid_id_set = set()
    with open(train_pth, "r") as f:
        for line in f:
            data = json.loads(line)
            valid_id_set.add(str(data['text_id']))
    return valid_id_set

def main():
    warnings.filterwarnings("ignore", category=UserWarning)
    parser = HfArgumentParser((CustomTrainingArguments, RunArguments))
    args, run_args = parser.parse_args_into_dataclasses()

    # We use wandb logger: https://wandb.ai/site.
    if args.local_rank == 0:  # only on main process
        # Initialize wandb run
        wandb.login()
        wandb.init(project="DSI2Table", name=args.run_name, entity="DSITable")

    tokenizer = MT5Tokenizer.from_pretrained(run_args.base_model_path, cache_dir='cache')
    fast_tokenizer = MT5TokenizerFast.from_pretrained(run_args.base_model_path, cache_dir='cache')

    # legal tokens
    SPIECE_UNDERLINE = "▁"
    INT_TOKEN_IDS = []
    for token, id in tokenizer.get_vocab().items():
        if token[0] == SPIECE_UNDERLINE:
            if token[1:].isdigit():
                INT_TOKEN_IDS.append(id)
        if token == SPIECE_UNDERLINE:
            INT_TOKEN_IDS.append(id)
        elif token.isdigit():
            INT_TOKEN_IDS.append(id)
    INT_TOKEN_IDS.append(tokenizer.eos_token_id)

    def restrict_decode_vocab(batch_idx, prefix_beam):
        return INT_TOKEN_IDS

    if run_args.task == "Index":
        training_args = args
        model = MT5ForConditionalGeneration.from_pretrained(run_args.base_model_path, cache_dir='cache')
        if run_args.peft:
            print("**************PEFT*****************")
            peft_config = LoraConfig(
                task_type=TaskType.SEQ_2_SEQ_LM,
                inference_mode=False,
                target_modules='(decoder\.block\.\d+\.layer\.2|encoder\.block\.\d+\.layer\.1)\.DenseReluDense\.(wi_0|wi_1|wo)',
                r=8,
                lora_alpha=32,
                lora_dropout=0.1
            )
            model = get_peft_model(model, peft_config)
            model.print_trainable_parameters()
        else:
            print("**************FULL*****************")

        train_dataset = IndexingTrainDataset(path_to_data=run_args.train_file,
                                             max_length=run_args.max_length,
                                             cache_dir='cache',
                                             tokenizer=tokenizer)

        valid_dataset = IndexingTrainDataset(path_to_data=run_args.valid_file,
                                             max_length=run_args.max_length,
                                             cache_dir='cache',
                                             remove_prompt=run_args.remove_prompt,
                                             tokenizer=tokenizer)
        all_dataset = IndexingTrainDataset(
            path_to_data=run_args.all_file,
            max_length=run_args.max_length,
            cache_dir='cache',
            tokenizer=tokenizer)

        trainer = DSITrainer(
            model=model,
            tokenizer=tokenizer,
            args=training_args,
            train_dataset=train_dataset,
            eval_dataset=valid_dataset,
            data_collator=IndexingCollator(
                tokenizer,
                padding='longest',
            ),
            compute_metrics = make_compute_metrics(fast_tokenizer, all_dataset.valid_ids),
            restrict_decode_vocab = restrict_decode_vocab,
            id_max_length=run_args.id_max_length,
        )
        trainer.train()


    elif run_args.task == 'Search':
        test_args = args
        test_args.do_train = False
        test_args.do_predict = True
        test_dataset = IndexingTrainDataset(path_to_data=run_args.valid_file,
                                            max_length=run_args.max_length,
                                            cache_dir='cache',
                                            remove_prompt=run_args.remove_prompt,
                                            tokenizer=tokenizer)
        test_q = []
        for data in test_dataset.train_data:
            test_q.append(str(data['text']))

        test_num = test_dataset.total_len
        print("test_num:{}".format(test_num))

        base_model = MT5ForConditionalGeneration.from_pretrained(run_args.base_model_path, cache_dir='cache')

        Lora_pths = [None, args.LoRA_1, args.LoRA_2, args.LoRA_3]
        partition_pths = [args.partition_0, args.partition_1, args.partition_2]
        Lora_pths = Lora_pths[:args.num]
        partition_pths = partition_pths[:args.num]
        cadidates_ids = [[] for _ in range(test_num)]
        '''
          candidates_ids = [ [ [model_0], [model_1]...], ...]
        '''
        lable_ids = []

        id2q = [dict() for _ in range(len(partition_pths))]
        '''
          id2q = [ {"text_id1": [q1, q2, ...]}, {"text_id2": [q3, q4, ...]} ]
        '''
        for i, (valid_pth, lora_pth) in enumerate(zip(partition_pths, Lora_pths)):
            valid_ids = valid_id(valid_pth)
            with open(valid_pth, 'r') as f:
                for line in f:
                    data = json.loads(line)
                    k = data["text_id"]
                    q = data["text"]
                    if k not in id2q[i].keys():
                        id2q[i][k] = []
                        # id2q[i][k].append(q)
                    else:
                        id2q[i][k].append(q)

            if lora_pth is not None:
                model = PeftModel.from_pretrained(base_model, lora_pth)
            else:

                model = base_model
            predictions, labels = inference_on_test_set(
                model=model,
                test_args=test_args,
                test_dataset=test_dataset,
                tokenizer=tokenizer,
                restrict_decode_vocab=restrict_decode_vocab,
                id_max_length=run_args.id_max_length
            )
            j = 0
            for beams, label in zip(predictions, labels):
                rank_list = tokenizer.batch_decode(beams, skip_special_tokens=True)
                label_id = tokenizer.decode(label, skip_special_tokens=True)
                lable_ids.append(label_id)
                filtered_ids = []
                for docid in rank_list:
                    if docid not in filtered_ids and docid in valid_ids:
                        filtered_ids.append(docid)
                cadidates_ids[j].append(filtered_ids[:5])
                j += 1
        hit_at_5 = 0
        hit_at_1 = 0
        model = SentenceTransformer("../model/stsb-roberta-base")
        merged_id2q = {}

        for d in id2q:
            merged_id2q.update(d)

        for i, q in enumerate(test_q):
            label_id = lable_ids[i]
            ids = list(itertools.chain(*cadidates_ids[i]))
            gen_qs = [merged_id2q[id] for id in ids]
            gen_qs = list(itertools.chain(*gen_qs))
            embed = model.encode(gen_qs, batch_size=2048)
            candidate_embeddings = np.array(embed)
            candidate_embeddings = candidate_embeddings / np.linalg.norm(candidate_embeddings, axis=1, keepdims=True)
            test_q_embeddings = model.encode(q)
            test_q_embeddings = test_q_embeddings.reshape(1, -1)
            test_q_embeddings = test_q_embeddings / np.linalg.norm(test_q_embeddings, axis=1, keepdims=True)
            max_ids = find_similar_embeddings(test_q_embeddings, candidate_embeddings, topk=5)
            t = v = 0
            id2model = dict()
            for k, sub_lst in enumerate(cadidates_ids[i]):
                for v in range(t, t + len(sub_lst) * 20):
                    id2model[v] = k
                t = v + 1
            maxids = max_ids.flatten().tolist()
            model_nums = [id2model[id] for id in maxids]
            count = Counter(model_nums)
            model_num = count.most_common(1)[0][0]
            preds = cadidates_ids[i][model_num]
            hits = np.where(np.array(preds)[:5] == label_id)[0]
            if len(hits) != 0:
                hit_at_5 += 1
                if hits[0] == 0:
                    hit_at_1 += 1
        print("Hits@1: {}".format(hit_at_1 / test_num))
        print("Hits@5: {}".format(hit_at_5 / test_num))

    else:
        raise NotImplementedError("--task should be in 'DSI' or 'inference'")


if __name__ == "__main__":

    main()