encode.py

import logging
import os
import pickle
import sys
from contextlib import nullcontext
from dataclasses import dataclass, field
import numpy as np
from tqdm import tqdm

import torch

from torch.utils.data import DataLoader
from transformers import AutoConfig, AutoTokenizer
from transformers import (
    HfArgumentParser,
)
from tevatron.arguments import ModelArguments, DataArguments, \
    TevatronTrainingArguments as TrainingArguments
from tevatron.data import EncodeDataset, EncodeCollator
from tevatron.modeling import EncoderOutput, DenseModel
from tevatron.datasets import HFQueryDataset, HFCorpusDataset

logger = logging.getLogger(__name__)


@dataclass
class xQGDataArguments(DataArguments):
    lang: str = field(
        default=None, metadata={"help": "Language of the generated queries."}
    )
    q_ind: int = field(
        default=0, metadata={"help": "Number of generated queries."}
    )


class xQGPreProcessor:
    def __init__(self, tokenizer, lang, q_ind, max_length=32):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.lang = lang
        self.q_ind = q_ind

    def __call__(self, example):
        docid = example['docid']
        generated_queries = example['generated_queries'][self.lang]
        query = self.tokenizer.encode(generated_queries[self.q_ind],
                                      add_special_tokens=False,
                                      max_length=self.max_length,
                                      truncation=True)
        return {'text_id': docid, 'text': query}


class xQGDataset(HFQueryDataset):
    def __init__(self, tokenizer, data_args, cache_dir):
        super().__init__(tokenizer, data_args, cache_dir)
        self.preprocessor = xQGPreProcessor
        self.data_args = data_args

    def process(self, shard_num=1, shard_idx=0):
        self.dataset = self.dataset.shard(shard_num, shard_idx)

        self.dataset = self.dataset.map(
            self.preprocessor(self.tokenizer, self.data_args.lang, self.data_args.q_ind, self.q_max_len),
            batched=False,
            num_proc=self.proc_num,
            remove_columns=self.dataset.column_names,
            desc=f"Running tokenization",
        )
        return self.dataset


def main():
    parser = HfArgumentParser((ModelArguments, xQGDataArguments, TrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
        model_args: ModelArguments
        data_args: xQGDataArguments
        training_args: TrainingArguments

    if training_args.local_rank > 0 or training_args.n_gpu > 1:
        raise NotImplementedError('Multi-GPU encoding is not supported.')

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,
    )

    num_labels = 1
    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        num_labels=num_labels,
        cache_dir=model_args.cache_dir,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir
    )

    model = DenseModel.load(
        model_name_or_path=model_args.model_name_or_path,
        config=config,
        cache_dir=model_args.cache_dir,
    )

    text_max_length = data_args.q_max_len
    encode_dataset = xQGDataset(tokenizer=tokenizer, data_args=data_args,
                                    cache_dir=data_args.data_cache_dir or model_args.cache_dir)

    encode_dataset = EncodeDataset(encode_dataset.process(data_args.encode_num_shard, data_args.encode_shard_index),
                                   tokenizer, max_len=text_max_length)

    encode_loader = DataLoader(
        encode_dataset,
        batch_size=training_args.per_device_eval_batch_size,
        collate_fn=EncodeCollator(
            tokenizer,
            max_length=text_max_length,
            padding='max_length'
        ),
        shuffle=False,
        drop_last=False,
        num_workers=training_args.dataloader_num_workers,
    )
    encoded = []
    lookup_indices = []
    model = model.to(training_args.device)
    model.eval()

    for (batch_ids, batch) in tqdm(encode_loader):
        lookup_indices.extend(batch_ids)
        with torch.cuda.amp.autocast() if training_args.fp16 else nullcontext():
            with torch.no_grad():
                for k, v in batch.items():
                    batch[k] = v.to(training_args.device)
                if data_args.encode_is_qry:
                    model_output: EncoderOutput = model(query=batch)
                    encoded.append(model_output.q_reps.cpu().detach().numpy())
                else:
                    model_output: EncoderOutput = model(passage=batch)
                    encoded.append(model_output.p_reps.cpu().detach().numpy())

    encoded = np.concatenate(encoded)


    with open(data_args.encoded_save_path, 'wb') as f:
        pickle.dump((encoded, lookup_indices), f)


if __name__ == "__main__":
    main()