run_ann_data_gen.py

# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys
import os
import numpy as np
from os.path import isfile, join
import argparse
import logging
import time
from functools import partial

import hnswlib
import paddle
from paddlenlp.transformers import AutoModel, AutoTokenizer
from paddlenlp.datasets import load_dataset, MapDataset, load_dataset
from paddlenlp.data import Stack, Tuple, Pad
from paddlenlp.utils.log import logger

from ance.model import SemanticIndexANCE
from data import get_latest_checkpoint, get_latest_ann_data
from data import convert_example, create_dataloader
from data import gen_id2corpus, gen_text_file
from ann_util import build_index

# yapf: disable
parser = argparse.ArgumentParser()

# Required parameters
parser.add_argument("--similar_text_pair_file", default=None, type=str,required=True, help="The train_set tsv file that each line is simialr text pair")
parser.add_argument("--corpus_file", default=None, type=str, required=True, help="The corpus file that each line is a text for buinding indexing")
parser.add_argument("--save_dir", default=None, type=str, required=True, help="Saved model dir, will look for latest checkpoint dir in here")
parser.add_argument("--ann_data_dir", default=None, type=str, required=True, help="The output directory where the training data will be written")

parser.add_argument("--init_from_ckpt", default=None, type=str, help="Initial model dir, will use this if no checkpoint is found in model_dir")
parser.add_argument("--end_ann_step", default=1000000, type=int, help="Stop after this number of data versions has been generated, default run forever")
parser.add_argument("--batch_size", default=128, type=int, help="Batch size for predicting embedding of texts")
parser.add_argument("--output_emb_size", default=None, type=int, help="output_embedding_size")

parser.add_argument("--max_seq_length", default=128, type=int, help="Batch size for predicting embedding of texts")
parser.add_argument("--topk_training", default=500, type=int, help="top k from which negative samples are collected")
parser.add_argument("--num_negative_sample", default=5, type=int, help="at each resample, how many negative samples per query do I use")

# hnsw argument
parser.add_argument("--hnsw_m", default=10, type=int, help="Recall number for each query from Ann index.")
parser.add_argument("--hnsw_ef", default=10, type=int, help="Recall number for each query from Ann index.")
parser.add_argument("--hnsw_max_elements", default=1000000, type=int, help="Recall number for each query from Ann index.")

args = parser.parse_args()
# yapf: enable


def generate_new_ann(args, data_loader_dict, checkpoint_path, latest_step_num):

    pretrained_model = AutoModel.from_pretrained("ernie-3.0-medium-zh")

    model = SemanticIndexANCE(pretrained_model, output_emb_size=args.output_emb_size)

    logger.info("checkpoint_path:{}".format(checkpoint_path))
    state_dict = paddle.load(checkpoint_path)

    model.set_dict(state_dict)
    logger.info("load params from:{}".format(checkpoint_path))

    logger.info("***** inference of corpus *****")
    final_index = build_index(args, data_loader_dict["corpus_data_loader"], model)

    logger.info("***** inference of query *****")
    query_embedding = model.get_semantic_embedding(data_loader_dict["text_data_loader"])

    text_list = data_loader_dict["text_list"]
    id2corpus = data_loader_dict["id2corpus"]
    text2similar_text = data_loader_dict["text2similar_text"]

    new_ann_data_path = os.path.join(args.ann_data_dir, str(latest_step_num))
    if not os.path.exists(new_ann_data_path):
        os.mkdir(new_ann_data_path)

    with open(os.path.join(new_ann_data_path, "new_ann_data"), "w") as f:
        for batch_index, batch_query_embedding in enumerate(query_embedding):
            recalled_idx, cosine_sims = final_index.knn_query(batch_query_embedding, args.topk_training)

            batch_size = len(cosine_sims)

            for row_index in range(batch_size):
                text_index = args.batch_size * batch_index + row_index

                hard_neg_samples = recalled_idx[row_index][-1 * args.num_negative_sample :]
                hard_neg_sims = cosine_sims[row_index][-1 * args.num_negative_sample :]

                for idx, hard_neg_doc_idx in enumerate(hard_neg_samples):
                    text = text_list[text_index]["text"]
                    similar_text = text2similar_text[text]
                    hard_neg_sample = id2corpus[hard_neg_doc_idx]
                    cosine_sim = 1.0 - hard_neg_sims[idx]
                    f.write("{}\t{}\t{}\n".format(text, similar_text, hard_neg_sample))

    succeed_flag_file = os.path.join(new_ann_data_path, "succeed_flag_file")
    open(succeed_flag_file, "a").close()
    logger.info("finish generate ann data step:{}".format(latest_step_num))


def build_data_loader(args, tokenizer):
    """build corpus_data_loader and text_data_loader"""

    id2corpus = gen_id2corpus(args.corpus_file)

    # conver_example function's input must be dict
    corpus_list = [{idx: text} for idx, text in id2corpus.items()]
    corpus_ds = MapDataset(corpus_list)

    trans_func = partial(convert_example, tokenizer=tokenizer, max_seq_length=args.max_seq_length)

    batchify_fn = lambda samples, fn=Tuple(
        Pad(axis=0, pad_val=tokenizer.pad_token_id),  # text_input
        Pad(axis=0, pad_val=tokenizer.pad_token_type_id),  # text_segment
    ): [data for data in fn(samples)]

    corpus_data_loader = create_dataloader(
        corpus_ds, mode="predict", batch_size=args.batch_size, batchify_fn=batchify_fn, trans_fn=trans_func
    )

    # build text data_loader
    text_list, text2similar_text = gen_text_file(args.similar_text_pair_file)

    text_ds = MapDataset(text_list)

    text_data_loader = create_dataloader(
        text_ds, mode="predict", batch_size=args.batch_size, batchify_fn=batchify_fn, trans_fn=trans_func
    )

    d = {
        "text_data_loader": text_data_loader,
        "corpus_data_loader": corpus_data_loader,
        "id2corpus": id2corpus,
        "text2similar_text": text2similar_text,
        "text_list": text_list,
    }

    return d


def ann_data_gen(args):
    # use init_from_ckpt as last_checkpoint
    last_checkpoint = args.init_from_ckpt

    # get latest_ann_data_step to decide when stop gen_ann_data
    _, latest_ann_data_step = get_latest_ann_data(args.ann_data_dir)

    rank = paddle.distributed.get_rank()
    if rank == 0:
        if not os.path.exists(args.ann_data_dir):
            os.makedirs(args.ann_data_dir)

    tokenizer = AutoTokenizer.from_pretrained("ernie-3.0-medium-zh")

    data_load_dict = build_data_loader(args, tokenizer)

    while latest_ann_data_step <= args.end_ann_step:
        next_checkpoint, latest_step_num = get_latest_checkpoint(args)
        logger.info("next_checkpoint:{}".format(next_checkpoint))

        if next_checkpoint == last_checkpoint:
            logger.info("next_checkpoint == lase_checkpoint:{}".format(next_checkpoint))
            logger.info("sleep 10s")
            time.sleep(10)
        else:
            logger.info("start generate ann data using checkpoint:{}".format(next_checkpoint))

            generate_new_ann(args, data_load_dict, next_checkpoint, latest_step_num)

            logger.info("finished generating ann data step {}".format(latest_step_num))

            last_checkpoint = next_checkpoint


def main():
    ann_data_gen(args)


if __name__ == "__main__":
    main()