/slm/pipelines/utils/offline_ann.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 argparse
import time

from pipelines.document_stores import (
    BaiduElasticsearchDocumentStore,
    ElasticsearchDocumentStore,
    MilvusDocumentStore,
)
from pipelines.nodes import DensePassageRetriever
from pipelines.utils import convert_files_to_dicts, fetch_archive_from_http, launch_es
from pipelines.utils.preprocessing import convert_files_to_dicts_splitter

data_dict = {
    "data/dureader_dev": "https://paddlenlp.bj.bcebos.com/applications/dureader_dev.zip",
    "data/baike": "https://paddlenlp.bj.bcebos.com/applications/baike.zip",
    "data/insurance": "https://paddlenlp.bj.bcebos.com/applications/insurance.zip",
    "data/file_example": "https://paddlenlp.bj.bcebos.com/pipelines/file_examples.zip",
}

# yapf: disable
parser = argparse.ArgumentParser()
parser.add_argument("--index_name", default="baike_cities", type=str, help="The index name of the ANN search engine")
parser.add_argument("--doc_dir", default="data/baike/", type=str, help="The doc path of the corpus")
parser.add_argument('--username', type=str, default="", help='Username of ANN search engine')
parser.add_argument('--password', type=str, default="", help='Password of ANN search engine')
parser.add_argument("--search_engine", choices=["elastic", "milvus", 'bes'], default="elastic", help="The type of ANN search engine.")
parser.add_argument("--host", type=str, default="127.0.0.1", help="host ip of ANN search engine")
parser.add_argument("--port", type=str, default="9200", help="port of ANN search engine")
parser.add_argument("--embedding_dim", default=768, type=int, help="The embedding_dim of index")
parser.add_argument("--split_answers", action="store_true", help="whether to split lines into question and answers")
parser.add_argument("--query_embedding_model", default="rocketqa-zh-base-query-encoder", type=str, help="The query_embedding_model path",)
parser.add_argument("--passage_embedding_model", default="rocketqa-zh-base-para-encoder", type=str, help="The passage_embedding_model path", )
parser.add_argument("--params_path", default="checkpoints/model_40/model_state.pdparams", type=str, help="The checkpoint path")
parser.add_argument("--delete_index", action="store_true", help="Whether to delete existing index while updating index")
parser.add_argument("--share_parameters", action="store_true", help="Use to control the query and title models sharing the same parameters",)
parser.add_argument('--model_type', choices=['ernie_search', 'ernie', 'bert', 'neural_search'], default="ernie", help="the ernie model types")
parser.add_argument('--embed_title', default=False, type=bool, help="The title to be  embedded into embedding")
parser.add_argument('--device', choices=['cpu', 'gpu'], default="gpu", help="Select devices, defaults to gpu.")
parser.add_argument('--search_fields', default=['content', 'name'], help="multi recall BM25Retriever set search_fields")
parser.add_argument('--use_splitter', default=False, type=bool, help="How to split documents")
parser.add_argument('--chunk_size', type=int, default=300, help="The length of data for indexing by retriever")
parser.add_argument('--chunk_overlap', type=int, default=0, help="a larger chunk than the chunk overlap")
parser.add_argument('--separator', type=str, default='\n', help="Use symbols to segment text, PDF, and image files, or connect some short chunks")
parser.add_argument('--filters', type=list, default=['\n'], help="Filter special symbols")
parser.add_argument('--language', type=str, default='chinese', help="the language of files")
parser.add_argument('--pooling_mode', choices=['max_tokens', 'mean_tokens', 'mean_sqrt_len_tokens', 'cls_token'], default='cls_token', help='the type of sentence embedding')
parser.add_argument("--es_chunk_size", default=500, type=int, help="Number of docs in one chunk sent to es")
parser.add_argument("--es_thread_count", default=32, type=int, help="Size of the threadpool to use for the bulk requests")
parser.add_argument("--es_queue_size", default=32, type=int, help="Size of the task queue between the main thread (producing chunks to send) and the processing threads.")
args = parser.parse_args()
# yapf: enable


def offline_ann(index_name, doc_dir):
    use_gpu = True if args.device == "gpu" else False
    if args.search_engine == "milvus":
        document_store = MilvusDocumentStore(
            embedding_dim=args.embedding_dim,
            host=args.host,
            index=args.index_name,
            port=args.port,
            index_param={"M": 16, "efConstruction": 50},
            index_type="HNSW",
        )
    elif args.search_engine == "bes":

        document_store = BaiduElasticsearchDocumentStore(
            host=args.host,
            port=args.port,
            username=args.username,
            password=args.password,
            embedding_dim=args.embedding_dim,
            similarity="dot_prod",
            vector_type="bpack_vector",
            search_fields=["content", "meta"],
            index=args.index_name,
            chunk_size=args.es_chunk_size,
            thread_count=args.es_thread_count,
            queue_size=args.es_queue_size,
        )

    else:
        launch_es()
        document_store = ElasticsearchDocumentStore(
            host=args.host,
            port=args.port,
            username=args.username,
            password=args.password,
            embedding_dim=args.embedding_dim,
            index=index_name,
            search_fields=args.search_fields,  # 当使用了多路召回并且搜索字段设置了除content的其他字段，构建索引时其他字段也需要设置，例如：['content', 'name']。
        )
    # 将每篇文档按照段落进行切分
    if args.use_splitter:
        dicts = convert_files_to_dicts_splitter(
            dir_path=doc_dir,
            split_paragraphs=True,
            split_answers=args.split_answers,
            encoding="utf-8",
            separator=args.separator,
            filters=args.filters,
            chunk_size=args.chunk_size,
            language=args.language,
            chunk_overlap=args.chunk_overlap,
        )
    else:
        dicts = convert_files_to_dicts(
            dir_path=doc_dir, split_paragraphs=True, split_answers=args.split_answers, encoding="utf-8"
        )

    print(dicts[:3])

    # 文档数据写入数据库
    document_store.write_documents(dicts)
    # 语义索引模型
    retriever = DensePassageRetriever(
        document_store=document_store,
        query_embedding_model=args.query_embedding_model,
        passage_embedding_model=args.passage_embedding_model,
        params_path=args.params_path,
        output_emb_size=args.embedding_dim if args.model_type in ["ernie_search", "neural_search"] else None,
        share_parameters=args.share_parameters,
        max_seq_len_query=64,
        max_seq_len_passage=256,
        batch_size=16,
        use_gpu=use_gpu,
        embed_title=args.embed_title,
    )
    # Writing docs may take a while. so waiting until writing docs to be completed.
    document_count = document_store.get_document_count()
    while document_count == 0:
        time.sleep(1)
        print("Waiting for writing docs to be completed.")
        document_count = document_store.get_document_count()
    # 建立索引库
    document_store.update_embeddings(retriever)


def delete_data(index_name):
    if args.search_engine == "milvus":
        document_store = MilvusDocumentStore(
            embedding_dim=args.embedding_dim,
            host=args.host,
            index=args.index_name,
            port=args.port,
            index_param={"M": 16, "efConstruction": 50},
            index_type="HNSW",
        )
    elif args.search_engine == "bes":

        document_store = BaiduElasticsearchDocumentStore(
            host=args.host,
            port=args.port,
            username=args.username,
            password=args.password,
            embedding_dim=args.embedding_dim,
            similarity="dot_prod",
            vector_type="bpack_vector",
            search_fields=["content", "meta"],
            index=args.index_name,
            chunk_size=args.es_chunk_size,
            thread_count=args.es_thread_count,
            queue_size=args.es_queue_size,
        )

    else:
        document_store = ElasticsearchDocumentStore(
            host=args.host,
            port=args.port,
            username="",
            password="",
            embedding_dim=args.embedding_dim,
            index=index_name,
        )
    document_store.delete_index(index_name)
    print("Delete an existing elasticsearch index {} Done.".format(index_name))


if __name__ == "__main__":
    if args.doc_dir in data_dict:
        fetch_archive_from_http(url=data_dict[args.doc_dir], output_dir=args.doc_dir)
    if args.delete_index:
        delete_data(args.index_name)
    offline_ann(args.index_name, args.doc_dir)