跳至内容

DAPR

论文解读

摘要

场景是从长文档(long document)中找到相关片段(passage),对SoTA分析发现错误主要源于错失了文档上下文内容,于是命名任务为DAPR——文档感知的片段检索(Document-Aware Passage Retrieval)。然后为此建立benchmark并进行实验。实验中对SoTA片段落检索器辅之以文档内容:(1)和BM25混合。(2)将文档信息掺和到段落信息里。实验结果:混合检索器法整体好,但是对于“需要文档级上下文信息”的问题完全失败;掺和法反之,整体较差但在上下文困难问题有进步。

介绍

IR有用。传统BM25通过统计方法。近期方法应用神经网络进行向量化表示实现复杂语义匹配。神经方法受限于“短片段”(short passage 如512tokens)输入。在现实中遇到长文本如学科论文时,超出限制而无法利用完整上下文。近期工作实现长文档encode并实现了文档级检索。然而返回长文档对于用户仍不便于获取有效信息,如在标注维基百科页面中与用户查询相关的段落时,35.8%的查询所对应相关段落平均位于文档第7.6段(标准差12.7)。

于是提出DAPR任务,就是在长文档中找到对应片段。举例解释“文档-感知” (Document-awared),如指代关系会导致Query的关键信息出现在不同片段里,可能导致关键信息所在片段遗漏。为深入了解,测试SoTA片段检索器并分析错误,发现53.5%的错误源于未结合文档上下文解析指代或主题,导致漏检正确答案段落。于是创建一个基准测试,其中包含来自不同领域的5个数据集,这些数据集提供了在关联文档中对相关段落的标注(就是问题有标出相关片段)。

实验中对SoTA片段落检索器辅之以文档内容:(1)和BM25混合。(2)将文档信息掺和到段落信息里。实验结果:混合检索器法整体好,但是对于“需要文档级上下文信息”的问题完全失败;掺和法反之,整体较差但在上下文困难问题有进步。这个基准测试对新的检索方法有意义。

相关工作

DocQA: 需要模型回答关于输入文档的一个问题,它的问题的相关文档是给定的,而DAPR需要从文档集里找相关信息(doc)。基于给定文档这个假设,DocQA的问题通常是和文档内容关联的,比如“哪个检索系统用于这个baseline”,这个问题在DocQA中因为给定了文档所以有默认关联,而转换到DAPR就不行了,“这个baseline"不知道是啥。

长文档检索:法1是MaxP,片段最高分作为文档分,法2是FirstP,第一片段得分作为文档分,法2不如法1。有学者比较了不同的“长距离注意力模型”的长文档检索。陈等人(2022年)提出了一种分层神经网络,并证明其优于基于MaxP。但这些工作相比DAPR,都没有研究 “考虑文档内容情况下的片段检索”。

混合检索:涉及多检索系统。其中排名混合将排名合并

实验复现代码

reproduce_bm25.sh

(.sh4种算法都类似,但passage_only.bm25大不相同)

export:将变量提升为 环境变量(全局变量),使得所有子进程(如其他脚本、Python程序等)都能继承这个变量。

直接设置:直接赋值定义的变量是 Shell变量(局部变量),仅在当前 Shell 会话中有效,子进程(如其他脚本、程序)无法访问。

1
2
3
4
5
6
7
# scripts/dgx2/exps/passage_only/bm25.sh
export NCCL_DEBUG="INFO"
export CUDA_VISIBLE_DEVICES="0"

# datasets=( "ConditionalQA" )
datasets=( "NaturalQuestions" )
# datasets=( "ConditionalQA" "MSMARCO" "NaturalQuestions" "Genomics" "MIRACL" )

for X in ${datasets[@]}: 遍历,取出每个元素到X

$ 变量替换:获取变量的值 ${变量名}

1
2
3
4
5
6
7
8
9
for dataset in ${datasets[@]}
do
    export DATA_DIR="data"
    export DATASET_PATH="$DATA_DIR/$dataset"
    export CLI_ARGS="
    --data_dir=$DATASET_PATH
    "
    # --data_dir=data/NaturalQuestions
    ...

命令替换:执行命令,并将输出结果作为字符串返回

$(python -m dapr.exps.passage_only.args.bm25 $CLI_ARGS):传入 CLI_ARGS 给dapr.exps.passage_only.args.bm25.py 并执行,捕获输出结果(如print的内容),py脚本代码如下,传入 –data_dir=data/NaturalQuestions,即py里有个变量data_dir=“data/NaturalQuestions”,BM25PassageOnlyArguments里的初始化会生成各种属性,包括output_dir。

1
2
3
4
if __name__ == "__main__":
    print(
        parse_cli(BM25PassageOnlyArguments).output_dir
    )  # For creating the logging path

export OUTPUT_DIR=$(…): … 即 bm25.py 里的输出被bash脚本捕获。

**mkdir -p $OUTPUT_DIR:**确保OUTPUT_DIR存在

**export LOG_PATH:**设置日志变量

1
2
3
4
5
6
7
	...
	export OUTPUT_DIR=$(python -m dapr.exps.passage_only.args.bm25 $CLI_ARGS)
	mkdir -p $OUTPUT_DIR
	export LOG_PATH="$OUTPUT_DIR/logging.log"
	echo "Logging file path: $LOG_PATH"
	...
	

**torchrun … > $LOG_PATH:**PyTorch 的分布式训练启动器,自动设置多进程环境,将命令的 标准输出(stdout) 重定向到文件,错误输出(stderr)仍会打印到终端。

  • –nproc_per_node=1:使用当前机器的1个GPU
  • –master_port=295001:主进程的通信端口
  • -m dapr.exps.passage_only.bm25:执行的脚本文件
  • $CLI_ARGS:传入的参数
1
2
3
	...
	torchrun --nproc_per_node=1 --master_port=29501 -m dapr.exps.passage_only.bm25 $CLI_ARGS > $LOG_PATH
done

dragon_plus (passage_only)

​ passage_only下的dragon_plus(类似上面的passage_only.bm25),即“仅段落检索使用dragon_plus检索器”,由以下 调用 实现。

​ reproduce_cp相当于main,其中主要执行了…passage_only.args.dragon_plus.py和…passage_only.dragon_plus_cp.py,前者用于配置一些参数,参数会传入后者,后者是实现功能和评估。

|——reproduce_cp.py
		|——dapr.exps.passage_only.args.dragon_plus.py 
		|				  |——dapr.exps.passage_only.args.base.py
		|								|——clddp.args.base
		|
		|——dapr.exps.passage_only.dragon_plus_cp.py
                            |——dapr.dataloader.py
                            |——dapr.retrievers.dense.py
                            |——dapr.dataloader.py
                            |——clddp.dm
                            |——clddp.search
                            |——clddp.evaluation
                            |——clddp.evaluation
                            |——clddp.dm

​ 后文中,以下简称,其他一致。

​ dapr.exps.passage_only.args.dragon_plus.py 简称为 …args.dragon_plus.py

​ dapr.exps.passage_only.dragon_plus_cp.py 简称为 …passage_only.dragon_plus_cp.py

reproduce_cp.py

|——reproduce_cp.py		替换原本的.sh,实现模拟命令行传参

​ reproduce_cp.py的功能为,替换原本的.sh,实现模拟命令行传参,而不用每次都在pycharm新建 运行/调试配置,大致流程和 reproduce_bm25.sh一样:

  • 设置全局/局部变量
  • 执行 dapr.exps.passage_only.args.dragon_plus,生成args
  • 再构造一些变量
  • 执行 dapr.exps.passage_only.dragon_plus_cp
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
import os
import subprocess
from pathlib import Path
import sys

def main():
    project_root = "/root/autodl-tmp/DAPR"  # 替换为你的项目路径
    os.environ["PYTHONPATH"] = f"{project_root}:{os.environ.get('PYTHONPATH', '')}"
    # 环境变量设置(对子进程生效)
    os.environ["NCCL_DEBUG"] = "INFO"
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    datasets = ["ConditionalQA"]
    data_dir = "data"

    for dataset in datasets:
        dataset_path = Path(data_dir) / dataset
        cli_args = f"--data_dir={dataset_path}"

        try:
            # 获取输出目录(需传递环境变量)
            output_dir = subprocess.check_output(
                [sys.executable, "-m", "dapr.exps.passage_only.args.dragon_plus", cli_args],
                text=True,
                env=os.environ  # 关键!确保子进程继承环境变量
            ).strip()

            # 创建目录
            output_path = Path(output_dir)
            output_path.mkdir(parents=True, exist_ok=True)

            # 日志路径
            log_path = output_path / "logging.log"
            print(f"Logging to: {log_path}")

            # 运行训练(合并 stdout 和 stderr)
            subprocess.run(
                [sys.executable, "-m", "dapr.exps.passage_only.dragon_plus_cp", cli_args],
                text=True,
                env=os.environ  # 关键!确保子进程继承环境变量
            )

        except subprocess.CalledProcessError as e:
            print(f"Error for {dataset}: {e}")
            continue

if __name__ == "__main__":
    main()

…args.dragon_plus.py

|——reproduce_cp.py		替换原本的.sh,实现模拟命令行传参
		|——...args.dragon_plus.py	定义类 BM25PassageOnlyArguments,输出output_dir

定义类 BM25PassageOnlyArguments 继承自 PassageOnlyArguments,重要属性如下:

  • run_name: str = XXX (自动生成的实验运行的唯一名称,见PassageOnlyArguments)
    
  • data_dir: str = os.path.join("exps/passage_only/bm25", self.run_name)     # 数据集路径
    
  • split: Split = Split.test               # 哪个数据部分(测试集 test 或训练集 train)
    
  • topk: int = 1000                        # 返回最相关的1000个结果
    
  • per_device_eval_batch_size: int = 32    # 每个GPU一次处理32条数据
    
  • fp16: bool = True                       # 用半精度浮点数(节省显存,加快计算)
    

不同点在于获取 parse_cli(BM25PassageOnlyArguments).output_dir 时使用的方法build_output_dir被重写为 os.path.join(“exps/passage_only/bm25”, self.run_name)

reproduce_cp执行该文件,进入if __name__ == “__main__",print(….output_dir),执行完reproduce_cp捕获该内容

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
from dataclasses import dataclass
import os
from clddp.utils import parse_cli
from dapr.exps.passage_only.args.base import PassageOnlyArguments


@dataclass
class BM25PassageOnlyArguments(PassageOnlyArguments):
    def build_output_dir(self) -> str:
        return os.path.join("exps/passage_only/bm25", self.run_name)


if __name__ == "__main__":
    print(
        parse_cli(BM25PassageOnlyArguments).output_dir
    )  # For creating the logging path

args.base.py

|——reproduce_cp.py		替换原本的.sh,实现模拟命令行传参
		|——...args.dragon_plus.py	定义类 BM25PassageOnlyArguments,输出output_dir
		|		|——args.base.py	定义类 PassageOnlyArguments和一些重要参数

定义类 PassageOnlyArguments,采用MixIn机制继承两个混入类(Mixin)(可插拔组件) ——AutoRunNameArgumentsMixIn, DumpableArgumentsMixIn,从clddp.args.base里import,继承前者能自动生成一次运行唯一名词(属性),继承后者获得输出目录和写入输出目录的方法。

设置了一些属性(见代码)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
from dataclasses import dataclass
from typing import Optional
from clddp.args.base import AutoRunNameArgumentsMixIn, DumpableArgumentsMixIn
from clddp.dm import Split

@dataclass
class PassageOnlyArguments(AutoRunNameArgumentsMixIn, DumpableArgumentsMixIn):
    data_dir: Optional[str] = None          # 数据集路径
    split: Split = Split.test               # 哪个数据部分(测试集 test 或训练集 train)
    topk: int = 1000                        # 返回最相关的1000个结果
    per_device_eval_batch_size: int = 32    # 每个GPU一次处理32条数据
    fp16: bool = True                       # 用半精度浮点数(节省显存,加快计算)

    def __post_init__(self) -> None:
        # 继承AutoRunNameArgumentsMixIn, DumpableArgumentsMixIn的初始化
        # 如self.output_dir = self.build_output_dir()    run_name: str = "auto_run_name"
        super().__post_init__()

clddp.args.base

定义两个类

  • AutoRunNameArgumentsMixIn:自动生成实验运行的唯一名称
  • DumpableArgumentsMixIn:设置output_dir和保存一些运行参数设置
|——reproduce_cp.py		替换原本的.sh,实现模拟命令行传参
		|——...args.dragon_plus.py	定义类 BM25PassageOnlyArguments,输出output_dir
		|	|——...args.base.py	定义类 PassageOnlyArguments和一些重要参数
		|				|——clddp.args.base  
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
from __future__ import annotations
from dataclasses import asdict, dataclass
import json
import logging
import os
from typing import List, Optional, Set, Type
from abc import ABC, abstractmethod

@dataclass
class AutoRunNameArgumentsMixIn:
    run_name: Optional[str] = None

    @property
    def escaped_args(self) -> Set[str]:
        """Escaped items for building the run_name."""
        return set()

    def get_arguments_from(self) -> List[type]:
        """Return the classes which the auto_run_name will get the arguments from."""
        return [self.__class__]

    @property
    def auto_run_name(self) -> str:
        arguments = [
            k
            for from_class in self.get_arguments_from()
            for k in from_class.__annotations__.keys()
            if k not in self.escaped_args
        ]
        items = [f"{argument}_{getattr(self, argument)}" for argument in arguments]
        # 用/连接每个"属性名_属性值"
        return "/".join(items)

    def __post_init__(self) -> None:
        if self.run_name is None:
            self.run_name = self.auto_run_name
        # 进行父类的__post_init__
        if hasattr(super(), "__post_init__"):
            super().__post_init__()  # Needed as a MixIn


@dataclass
class DumpableArgumentsMixIn(ABC):
    output_dir: Optional[str] = None

    def __post_init__(self) -> None:
        if self.output_dir is None:
            self.output_dir = self.build_output_dir()
        if hasattr(super(), "__post_init__"):
            super().__post_init__()  # Needed as a MixIn

    @abstractmethod
    def build_output_dir(self) -> str:
        pass

    def dump_arguments(self, fname: str = "arguments.json") -> None:
        os.makedirs(self.output_dir, exist_ok=True)
        fargs = os.path.join(self.output_dir, fname)
        with open(fargs, "w") as f:
            json.dump(asdict(self), f, indent=4)
        logging.info(f"Dumped arguments to {fargs}")

…passage_only.dragon_plus_cp.py

|——reproduce_cp.py		替换原本的.sh,实现模拟命令行传参
		|——...args.dragon_plus.py	定义类 BM25PassageOnlyArguments,输出output_dir
		|	|——...args.base.py	定义类 PassageOnlyArguments和一些重要参数
		|				|——clddp.args.base  定义两个MixIn
---------|----------------------------------------------------------后面省略上面内容
		|——...passage_only.dragon_plus_cp.py

其代码单步执行会顺序进出如下模块。

|——...args.dragon_plus			加载参数和保存参数到json文件
|——clddp.utils				    初始化ddp
|——dapr.dataloader.py			DAPRDataConfig, DAPRDataLoader实现数据加载
|——dapr.retrievers.dense.py		引入检索器
|——dapr.dataloader.py			获取标注好的问题
|——clddp.dm						
|——clddp.search
|——clddp.evaluation
|——clddp.evaluation
|——clddp.dm

代码

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
import json
import logging
import os
from clddp.dm import LabeledQuery, RetrievedPassageIDList
from clddp.utils import initialize_ddp, is_device_zero, parse_cli, set_logger_format
from dapr.retrievers.dense import DRAGONPlus
from dapr.exps.passage_only.args.dragon_plus import DRAGONPlusPassageOnlyArguments
from clddp.search import search
from dapr.dataloader import DAPRDataConfig, DAPRDataLoader
from clddp.evaluation import RetrievalEvaluator
import torch

def is_device_zero():
    return True
def main():
    set_logger_format(logging.INFO if is_device_zero() else logging.WARNING)
    
    # ...args.dragon_plus ===============================
    args = parse_cli(DRAGONPlusPassageOnlyArguments)  # 加载参数
    args.dump_arguments()	# 记录参数到arguments.json
    
    # 设置环境变量 ===============================
    if not torch.distributed.is_initialized():
        os.environ["MASTER_ADDR"] = "localhost"
        os.environ["MASTER_PORT"] = "29501"
        os.environ["WORLD_SIZE"] = "1"
        os.environ["RANK"] = "0"
        os.environ["LOCAL_RANK"] = "0"
        torch.distributed.init_process_group(backend="nccl")
	
    # clddp.utils ===============================
    initialize_ddp() 	# 初始化 dstributed data process

    # dapr.dataloader ===============================
    dataset = DAPRDataLoader(DAPRDataConfig(args.data_dir)).load_data(
        progress_bar=is_device_zero()
    )

    # dapr.retrievers.dense ================================
    retriever = DRAGONPlus()

    # dapr.dataloader ===============================
    labeled_queries = dataset.get_labeled_queries(args.split)

    # clddp.dm ==================================
    queries = LabeledQuery.get_unique_queries(labeled_queries)

    # clddp.search
    retrieved = search(
        retriever=retriever,
        collection_iter=dataset.collection_iter,
        collection_size=dataset.collection_size,
        queries=queries,
        topk=args.topk,
        per_device_eval_batch_size=args.per_device_eval_batch_size,
        fp16=args.fp16,
    )

    # clddp.evaluation
    evaluator = RetrievalEvaluator(eval_dataset=dataset, split=args.split)
    report = evaluator(retrieved)

    freport = os.path.join(args.output_dir, "metrics.json")
    with open(freport, "w") as f:
        json.dump(report, f, indent=4)
    logging.info(f"Saved evaluation metrics to {freport}.")
    franked = os.path.join(args.output_dir, "ranking_results.txt")
    RetrievedPassageIDList.dump_trec_csv(
        retrieval_results=retrieved, fpath=franked, system="dragon_plus"
    )
    logging.info(f"Saved ranking results to {franked}.")

if __name__ == "__main__":
    main()

clddp.utils

----|---------------------上面省略
    |——...passage_only.dragon_plus_cp.py
            |——...args.dragon_plus		加载参数和保存参数到json文件
            |——clddp.utils				初始化ddp

dapr.dataloader.py

----|---------------------上面省略
    |——...passage_only.dragon_plus_cp.py
            |——...args.dragon_plus		加载参数和保存参数到json文件
            |——clddp.utils				初始化ddp
            |——dapr.dataloader.py		DAPRDataConfig, DAPRDataLoader实现数据加载

dragon_plus_cp中语句:

1
2
3
4
# dapr.dataloader ===============================
    dataset = DAPRDataLoader(DAPRDataConfig(args.data_dir)).load_data(
        progress_bar=is_device_zero()
    )

分两步:

  • DAPRDataConfig(args.data_dir)

  • DAPRDataLoader(…).load_data( progress_bar=is_device_zero() )

    where is_device_zero()==True 单卡模式

1
2
3
4
5
6
7
8
# DAPRDataConfig(args.data_dir)
@dataclass
class DAPRDataConfig:
    data_name_or_path: str # 路径字符串"os.path.join("exps/passage_only/dragon_plus", self.run_name)"
    titled: bool = False
    # 类RetrievalLevel和ParagraphSeparator其实就是C中结构体,静态的
    retrieval_level: RetrievalLevel = RetrievalLevel.paragraph  # "paragraph"
    paragraph_separator: ParagraphSeparator = ParagraphSeparator.blank  # "blank"

(DAPRDataLoader只列出用到的内容)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class DAPRDataLoader:
    def __init__(self, config: DAPRDataConfig) -> None:
        self.config = config
	...
    def load_data(self, progress_bar: bool) -> RetrievalDataset:
        data = LoadedData.from_dump(self.config.data_name_or_path, pbar=progress_bar)
        assert data.meta_data is not None
        
        retrieval_level = self.config.retrieval_level
        if retrieval_level is RetrievalLevel.document:
            collection_size = data.meta_data["ndocs"]
        else:
            collection_size = (
                data.meta_data["nchunks"]
                if data.meta_data.get("nchunks_candidates") is None
                else data.meta_data["nchunks_candidates"]
            )
        dataset = RetrievalDataset(
            collection_iter_fn=partial(self.collection_iter_fn, data=data),
            collection_size=collection_size,
            train_labeled_queries=self.build_labeled_queries(
                labeled_queries=data.labeled_queries_train
            ),
            dev_labeled_queries=self.build_labeled_queries(
                labeled_queries=data.labeled_queries_dev
            ),
            test_labeled_queries=self.build_labeled_queries(
                labeled_queries=data.labeled_queries_test
            ),
        )
        return dataset    

dapr.datasets.dm

----|---------------------上面省略
    |——...passage_only.dragon_plus_cp.py
            |——...args.dragon_plus		加载参数和保存参数到json文件
            |——clddp.utils				初始化ddp
            |——dapr.dataloader.py		DAPRDataConfig, DAPRDataLoader实现数据加载
            	|——dapr.datasets.dm		
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
@dataclass
class LoadedData:
    # data:
    corpus_iter_fn: Optional[Callable[[], Iterable[Document]]] = None
    labeled_queries_train: Optional[List[LabeledQuery]] = None
    labeled_queries_dev: Optional[List[LabeledQuery]] = None
    labeled_queries_test: Optional[List[LabeledQuery]] = None

    # meta data:
    meta_data: Optional[MetaDataJson] = None
    
    ...
    @staticmethod
    def build_corpus_iter_fn(fpath: str) -> Iterable[Document]:
        with open(fpath) as f:
            for line in f:
                doc_json: DocumentJson = ujson.loads(line)
                document = Document.from_json(doc_json)
                document.set_default_candidates()  # Also for compatibility with the older version
                yield document
                
    @staticmethod
    def load_labeled_queries(
        fpath: str, total: int, pbar: bool
    ) -> Optional[List[LabeledQuery]]:
        # fpath 文件路径 total 问答对个数 pbar对于device 0 是true
        if not os.path.exists(fpath):
            return None

        labeled_queries: List[LabeledQuery] = []
        with open(fpath) as f:
            for line in tqdm.tqdm(
                f, total=total, desc=f"Loading from {fpath}", disable=not pbar
            ):
                lq_json: LabeledQueryJson = ujson.loads(line)
                query_json = lq_json["query"]
                query = Query(query_id=query_json["id"], text=query_json["text"])
                judged_chunk_jsons = lq_json["judged_chunks"]
                judged_chunks = []
                for jchk_json in judged_chunk_jsons:
                    document = Document.from_json(jchk_json["belonging_doc"])
                    chunk_json: ChunkJson = jchk_json["chunk"]
                    chunk = Chunk(
                        chunk_id=chunk_json["id"],
                        text=chunk_json["text"],
                        doc_summary=None,
                        belonging_doc=document,
                    )
                    jchk = JudgedChunk(
                        query=query, chunk=chunk, judgement=jchk_json["judgement"]
                    )
                    judged_chunks.append(jchk)
                labeled_query = LabeledQuery(query=query, judged_chunks=judged_chunks)
                labeled_queries.append(labeled_query)
        return labeled_queries

    
    @classmethod
    def from_dump(
        cls: Type[LoadedData], dump_dir: str, pbar: bool = True
    ) -> LoadedData:
        # dump_dir=".../ConditionalQA"
        # 加载meta_data.json里的数据
        with open(os.path.join(dump_dir, "meta_data.json")) as f:
            meta_data: MetaDataJson = ujson.load(f)
            meta_data["chunk_separator"] = Separator(meta_data["chunk_separator"])
        # 实例化并传出
        loaded_data = cls(
            corpus_iter_fn=partial(
                cls.build_corpus_iter_fn, os.path.join(dump_dir, "corpus.jsonl")
            ),
            labeled_queries_train=cls.load_labeled_queries(
                fpath=os.path.join(dump_dir, "train.jsonl"),
                total=meta_data["nqueries_train"],
                pbar=pbar,
            ),
            labeled_queries_dev=cls.load_labeled_queries(
                fpath=os.path.join(dump_dir, "dev.jsonl"),
                total=meta_data["nqueries_dev"],
                pbar=pbar,
            ),
            labeled_queries_test=cls.load_labeled_queries(
                fpath=os.path.join(dump_dir, "test.jsonl"),
                total=meta_data["nqueries_test"],
                pbar=pbar,
            ),
            meta_data=meta_data,
        )

        return loaded_data

总结

屎一样的代码

最后更新于