flashrag-dev 0.1.4.dev20250310__tar.gz → 0.1.4.dev20250321__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (62) hide show
  1. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/PKG-INFO +29 -21
  2. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/README.md +28 -20
  3. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/config/basic_config.yaml +2 -0
  4. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/dataset/dataset.py +7 -0
  5. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/generator/openai_generator.py +1 -1
  6. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/pipeline/__init__.py +2 -1
  7. flashrag_dev-0.1.4.dev20250321/flashrag/pipeline/reasoning_pipeline.py +154 -0
  8. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/retriever/encoder.py +8 -6
  9. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/retriever/retriever.py +97 -63
  10. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/retriever/utils.py +11 -4
  11. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/utils/utils.py +8 -0
  12. flashrag_dev-0.1.4.dev20250321/flashrag/version.py +1 -0
  13. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag_dev.egg-info/PKG-INFO +29 -21
  14. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag_dev.egg-info/SOURCES.txt +1 -0
  15. flashrag_dev-0.1.4.dev20250310/flashrag/version.py +0 -1
  16. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/LICENSE +0 -0
  17. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/__init__.py +0 -0
  18. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/config/__init__.py +0 -0
  19. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/config/config.py +0 -0
  20. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/dataset/__init__.py +0 -0
  21. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/dataset/utils.py +0 -0
  22. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/evaluator/__init__.py +0 -0
  23. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/evaluator/_bleu.py +0 -0
  24. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/evaluator/evaluator.py +0 -0
  25. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/evaluator/metrics.py +0 -0
  26. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/evaluator/utils.py +0 -0
  27. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/generator/__init__.py +0 -0
  28. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/generator/fid.py +0 -0
  29. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/generator/generator.py +0 -0
  30. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/generator/multimodal_generator.py +0 -0
  31. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/generator/stop_word_criteria.py +0 -0
  32. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/generator/utils.py +0 -0
  33. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/judger/__init__.py +0 -0
  34. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/judger/judger.py +0 -0
  35. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/pipeline/active_pipeline.py +0 -0
  36. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/pipeline/branching_pipeline.py +0 -0
  37. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/pipeline/mm_pipeline.py +0 -0
  38. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/pipeline/pipeline.py +0 -0
  39. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/pipeline/replug_utils.py +0 -0
  40. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/prompt/__init__.py +0 -0
  41. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/prompt/base_prompt.py +0 -0
  42. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/prompt/mm_prompt.py +0 -0
  43. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/prompt/selfask_examplars.py +0 -0
  44. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/prompt/trace_examplars.py +0 -0
  45. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/refiner/__init__.py +0 -0
  46. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/refiner/kg_refiner.py +0 -0
  47. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/refiner/llmlingua_compressor.py +0 -0
  48. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/refiner/refiner.py +0 -0
  49. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/refiner/selective_context_compressor.py +0 -0
  50. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/retriever/__init__.py +0 -0
  51. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/retriever/__main__.py +0 -0
  52. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/retriever/index_builder.py +0 -0
  53. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/retriever/reranker.py +0 -0
  54. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/utils/__init__.py +0 -0
  55. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/utils/constants.py +0 -0
  56. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag/utils/pred_parse.py +0 -0
  57. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag_dev.egg-info/dependency_links.txt +0 -0
  58. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag_dev.egg-info/requires.txt +0 -0
  59. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/flashrag_dev.egg-info/top_level.txt +0 -0
  60. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/pyproject.toml +0 -0
  61. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/setup.cfg +0 -0
  62. {flashrag_dev-0.1.4.dev20250310 → flashrag_dev-0.1.4.dev20250321}/setup.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: flashrag_dev
3
- Version: 0.1.4.dev20250310
3
+ Version: 0.1.4.dev20250321
4
4
  Summary: A library for efficient Retrieval-Augmented Generation research
5
5
  Home-page: https://github.com/RUC-NLPIR/FlashRAG
6
6
  Author: Jiajie Jin, Yutao Zhu, Chenghao Zhang, Xinyu Yang, Zhicheng Dou
@@ -41,7 +41,7 @@ License-File: LICENSE
41
41
  </h4>
42
42
 
43
43
 
44
- FlashRAG is a Python toolkit for the reproduction and development of Retrieval Augmented Generation (RAG) research. Our toolkit includes 36 pre-processed benchmark RAG datasets and 16 state-of-the-art RAG algorithms.
44
+ FlashRAG is a Python toolkit for the reproduction and development of Retrieval Augmented Generation (RAG) research. Our toolkit includes 36 pre-processed benchmark RAG datasets and 17 state-of-the-art RAG algorithms.
45
45
 
46
46
  <p align="center">
47
47
  <img src="asset/framework.jpg">
@@ -75,7 +75,7 @@ https://github.com/user-attachments/assets/8ca00873-5df2-48a7-b853-89e7b18bc6e9
75
75
 
76
76
  - **Comprehensive Benchmark Datasets**: A collection of 36 pre-processed RAG benchmark datasets to test and validate RAG models' performances.
77
77
 
78
- - **Pre-implemented Advanced RAG Algorithms**: Features 16 advancing RAG algorithms with reported results, based on our framework. Easily reproducing results under different settings.
78
+ - **Pre-implemented Advanced RAG Algorithms**: Features 17 advancing RAG algorithms with reported results, based on our framework. Easily reproducing results under different settings.
79
79
 
80
80
  - **Efficient Preprocessing Stage**: Simplifies the RAG workflow preparation by providing various scripts like corpus processing for retrieval, retrieval index building, and pre-retrieval of documents.
81
81
 
@@ -96,17 +96,19 @@ FlashRAG is still under development and there are many issues and room for impro
96
96
  - [ ] Enhance code adaptability and readability
97
97
 
98
98
  ## :page_with_curl: Changelog
99
+ [25/03/21] We have added support for **Reasoning Pipeline**, which is a new paradigm that combines reasoning ability and retrieval, representing work that includes [Search-o1](https://github.com/sunnynexus/Search-o1), [R1-Searcher](https://github.com/SsmallSong/R1-Searcher), [ReSearch](https://github.com/Agent-RL/ReSearch). We evaluate the performance of the pipeline on various RAG benchmarks, it can achieve F1 scores close to 60 on multi hop inference datasets such as Hotpotqe. See it in [**result table**](#robot-supporting-methods).
100
+
99
101
  [25/02/24] 🔥🔥🔥 We have added support for **multimodal RAG**, including [**MLLMs like Llava, Qwen, InternVL**](https://ruc-nlpir.github.io/FlashRAG/#/zh-cn/component/generator?id=%e5%a4%9a%e6%a8%a1%e6%80%81%e7%94%9f%e6%88%90%e5%99%a8), and various [**multimodal retrievers with Clip architecture**](https://ruc-nlpir.github.io/FlashRAG/#/zh-cn/component/retriever?id=%e5%a4%9a%e6%a8%a1%e6%80%81%e6%a3%80%e7%b4%a2%e5%99%a8). More information can be found in our new version of arxiv article and our documentation. Try it!
100
102
 
101
103
  [25/01/21] Our technical paper [FlashRAG: A Python Toolkit for Efficient RAG Research](https://arxiv.org/abs/2405.13576) is honored to have been accepted to the Resource Track of the 2025 **ACM Web Conference (WWW 2025)**. Please Check it out!
102
104
 
103
105
  [25/01/12] Introduce <strong>FlashRAG-UI</strong>, an easy to use interface. You can easily and quickly configure and experience the supported RAG methods and evaluate them on the benchmarks.
104
106
 
105
- [25/01/11] We have added support for a new method [<u>RQRAG</u>](https://arxiv.org/abs/2404.00610) method, see it in [**reproduce_experiment**](docs/reproduce_experiment.md).
107
+ [25/01/11] We have added support for a new method [<u>RQRAG</u>](https://arxiv.org/abs/2404.00610) method, see it in [**reproduce_experiment**](docs/original_docs/reproduce_experiment.md).
106
108
 
107
- [25/01/07] We have currently support the aggregation of multiple retrievers, see it in [**multi retriever usage**](https://github.com/RUC-NLPIR/FlashRAG/blob/main/docs/multi_retriever_usage.md).
109
+ [25/01/07] We have currently support the aggregation of multiple retrievers, see it in [**multi retriever usage**](https://github.com/RUC-NLPIR/FlashRAG/blob/main/docs/original_docs/multi_retriever_usage.md).
108
110
 
109
- [25/01/07] We have integrated a very flexible and lightweight corpus chunking library [**Chunkie**](https://github.com/chonkie-ai/chonkie?tab=readme-ov-file#usage), which supports various custom chunking methods (tokens, sentences, semantic, etc.). Use it in [<u>chunking doc corpus</u>](docs/chunk-doc-corpus.md).
111
+ [25/01/07] We have integrated a very flexible and lightweight corpus chunking library [**Chunkie**](https://github.com/chonkie-ai/chonkie?tab=readme-ov-file#usage), which supports various custom chunking methods (tokens, sentences, semantic, etc.). Use it in [<u>chunking doc corpus</u>](docs/original_docs/chunk-doc-corpus.md).
110
112
 
111
113
  [24/10/21] We have released a version based on the Paddle framework that supports Chinese hardware platforms. Please refer to [FlashRAG Paddle](https://github.com/RUC-NLPIR/FlashRAG-Paddle) for details.
112
114
 
@@ -125,7 +127,7 @@ FlashRAG is still under development and there are many issues and room for impro
125
127
 
126
128
  [24/07/17] Due to some unknown issues with HuggingFace, our original dataset link has been invalid. We have updated it. Please check the [new link](https://huggingface.co/datasets/RUC-NLPIR/FlashRAG_datasets/) if you encounter any problems.
127
129
 
128
- [24/07/06] We add support for a new method: [<u>Trace</u>](https://arxiv.org/abs/2406.11460), which refine text by constructing a knowledge graph. See it [<u>results</u>](#robot-supporting-methods) and [<u>details</u>](./docs/baseline_details.md).
130
+ [24/07/06] We add support for a new method: [<u>Trace</u>](https://arxiv.org/abs/2406.11460), which refine text by constructing a knowledge graph. See it [<u>results</u>](#robot-supporting-methods) and [<u>details</u>](./docs/original_docs/baseline_details.md).
129
131
 
130
132
  [24/06/19] We add support for a new method: [<u>IRCoT</u>](https://arxiv.org/abs/2212.10509), and update the [<u>result table</u>](#robot-supporting-methods).
131
133
 
@@ -133,9 +135,9 @@ FlashRAG is still under development and there are many issues and room for impro
133
135
 
134
136
  [24/06/11] We have integrated `sentence transformers` in the retriever module. Now it's easier to use the retriever without setting pooling methods.
135
137
 
136
- [24/06/05] We have provided detailed document for reproducing existing methods (see [how to reproduce](./docs/reproduce_experiment.md), [baseline details](./docs/baseline_details.md)), and [<u>configurations settings</u>](./docs/configuration.md).
138
+ [24/06/05] We have provided detailed document for reproducing existing methods (see [how to reproduce](./docs/original_docs/reproduce_experiment.md), [baseline details](./docs/original_docs/baseline_details.md)), and [<u>configurations settings</u>](./docs/original_docs/configuration.md).
137
139
 
138
- [24/06/02] We have provided an introduction of FlashRAG for beginners, see [<u>an introduction to flashrag</u>](./docs/introduction_for_beginners_en.md) ([<u>中文版</u>](./docs/introduction_for_beginners_zh.md) [<u>한국어</u>](./docs/introduction_for_beginners_kr.md)).
140
+ [24/06/02] We have provided an introduction of FlashRAG for beginners, see [<u>an introduction to flashrag</u>](./docs/original_docs/introduction_for_beginners_en.md) ([<u>中文版</u>](./docs/original_docs/introduction_for_beginners_zh.md) [<u>한국어</u>](./docs/original_docs/introduction_for_beginners_kr.md)).
139
141
 
140
142
  [24/05/31] We supported Openai-series models as generator.
141
143
 
@@ -200,7 +202,7 @@ To build an index, you first need to save your corpus as a `jsonl` file with eac
200
202
  {"id": "1", "contents": "..."}
201
203
  ```
202
204
 
203
- If you want to use Wikipedia as your corpus, you can refer to our documentation [Processing Wikipedia](./docs/process-wiki.md) to convert it into an indexable format.
205
+ If you want to use Wikipedia as your corpus, you can refer to our documentation [Processing Wikipedia](./docs/original_docs/process-wiki.md) to convert it into an indexable format.
204
206
 
205
207
  ### Index Construction
206
208
 
@@ -289,7 +291,7 @@ my_config = Config(
289
291
  config_dict = config_dict
290
292
  ```
291
293
 
292
- We provide comprehensive guidance on how to set configurations, you can see our [<u>configuration guidance</u>](./docs/configuration.md).
294
+ We provide comprehensive guidance on how to set configurations, you can see our [<u>configuration guidance</u>](./docs/original_docs/configuration.md).
293
295
  You can also refer to the [<u>basic yaml file</u>](./flashrag/config/basic_config.yaml) we provide to set your own parameters.
294
296
 
295
297
  Next, load the corresponding dataset and initialize the pipeline. The components in the pipeline will be automatically loaded.
@@ -360,11 +362,11 @@ class ToyPipeline(BasicPipeline):
360
362
  return dataset
361
363
  ```
362
364
 
363
- Please first understand the input and output forms of the components you need to use from our [<u>documentation</u>](./docs/basic_usage.md).
365
+ Please first understand the input and output forms of the components you need to use from our [<u>documentation</u>](./docs/original_docs/basic_usage.md).
364
366
 
365
367
  ### Just use components
366
368
 
367
- If you already have your own code and only want to use our components to embed the original code, you can refer to the [<u>basic introduction of the components</u>](./docs/basic_usage.md) to obtain the input and output formats of each component.
369
+ If you already have your own code and only want to use our components to embed the original code, you can refer to the [<u>basic introduction of the components</u>](./docs/original_docs/basic_usage.md) to obtain the input and output formats of each component.
368
370
 
369
371
  ## :gear: Components
370
372
 
@@ -482,7 +484,7 @@ In each category, we have implemented corresponding common pipelines. Some pipel
482
484
  <td>Ranking and merging generated results based on each document</td>
483
485
  </tr>
484
486
  <tr>
485
- <td rowspan="5">Loop</td>
487
+ <td rowspan="6">Loop</td>
486
488
  <td>Iterative Pipeline</td>
487
489
  <td>Alternating retrieval and generation</td>
488
490
  </tr>
@@ -502,6 +504,10 @@ In each category, we have implemented corresponding common pipelines. Some pipel
502
504
  <td>IRCoT Pipeline</td>
503
505
  <td>Integrate retrieval process with CoT</td>
504
506
  </tr>
507
+ <tr>
508
+ <td>Reasoning Pipeline</td>
509
+ <td>Reasoning with retrieval</td>
510
+ </tr>
505
511
  </tbody>
506
512
  </table>
507
513
 
@@ -554,11 +560,11 @@ We have implemented 15 works with a consistent setting of:
554
560
 
555
561
  - **Generator:** LLAMA3-8B-instruct with input length of 2048
556
562
  - **Retriever:** e5-base-v2 as embedding model, retrieve 5 docs per query
557
- - **Prompt:** A consistent default prompt, template can be found in the [<u>method details</u>](./docs/baseline_details.md).
563
+ - **Prompt:** A consistent default prompt, template can be found in the [<u>method details</u>](./docs/original_docs/baseline_details.md).
558
564
 
559
565
  For open-source methods, we implemented their processes using our framework. For methods where the author did not provide source code, we will try our best to follow the methods in the original paper for implementation.
560
566
 
561
- For necessary settings and hyperparameters specific to some methods, we have documented them in the **specific settings** column. For more details, please consult our [<u>reproduce guidance</u>](./docs/reproduce_experiment.md) and [<u>method details</u>](./docs/baseline_details.md).
567
+ For necessary settings and hyperparameters specific to some methods, we have documented them in the **specific settings** column. For more details, please consult our [<u>reproduce guidance</u>](./docs/original_docs/reproduce_experiment.md) and [<u>method details</u>](./docs/original_docs/baseline_details.md).
562
568
 
563
569
  It’s important to note that, to ensure consistency, we have utilized a uniform setting. However, this setting may differ from the original setting of the method, leading to variations in results compared to the original outcomes.
564
570
 
@@ -582,6 +588,8 @@ It’s important to note that, to ensure consistency, we have utilized a uniform
582
588
  | [Iter-Retgen](https://arxiv.org/abs/2305.15294), [ITRG](https://arxiv.org/abs/2310.05149) | Loop | 36.8 | 60.1 | 38.3 | 21.6 | 37.9 | 18.2 | |
583
589
  | [IRCoT](https://aclanthology.org/2023.acl-long.557.pdf) | Loop | 33.3 | 56.9 | 41.5 | 32.4 | 45.6 | 20.7 | |
584
590
  | [RQRAG](https://arxiv.org/abs/2404.00610) | Loop | 32.6 | 52.5 | 33.5 | 35.8 | 46.4 | 26.2 | Use trained rqrag-llama2-7B |
591
+ | [R1-Searcher](https://arxiv.org/pdf/2503.05592) | Reasoning | 37.3 | 58.7 | 59.5 | 55.5 | 43.9 | 27.5 | Using trained Qwen2.5-7B model |
592
+
585
593
 
586
594
 
587
595
  ## :notebook: Supporting Datasets & Document Corpus
@@ -653,7 +661,7 @@ Our toolkit supports jsonl format for retrieval document collections, with the f
653
661
 
654
662
  The `contents` key is essential for building the index. For documents that include both text and title, we recommend setting the value of `contents` to `{title}\n{text}`. The corpus file can also contain other keys to record additional characteristics of the documents.
655
663
 
656
- In the academic research, Wikipedia and MS MARCO are the most commonly used retrieval document collections. For Wikipedia, we provide a [<u>comprehensive script</u>](./docs/process-wiki.md) to process any Wikipedia dump into a clean corpus. Additionally, various processed versions of the Wikipedia corpus are available in many works, and we have listed some reference links.
664
+ In the academic research, Wikipedia and MS MARCO are the most commonly used retrieval document collections. For Wikipedia, we provide a [<u>comprehensive script</u>](./docs/original_docs/process-wiki.md) to process any Wikipedia dump into a clean corpus. Additionally, various processed versions of the Wikipedia corpus are available in many works, and we have listed some reference links.
657
665
 
658
666
  For MS MARCO, it is already processed upon release and can be directly downloaded from its [<u>hosting link</u>](https://huggingface.co/datasets/Tevatron/msmarco-passage-corpus) on Hugging Face.
659
667
 
@@ -671,10 +679,10 @@ The index was created using the e5-base-v2 retriever on our uploaded wiki18_100w
671
679
 
672
680
  ## :raised_hands: Additional FAQs
673
681
 
674
- - [How should I set different experimental parameters?](./docs/configuration.md)
675
- - [How to build my own corpus, such as a specific segmented Wikipedia?](./docs/process-wiki.md)
676
- - [How to index my own corpus?](./docs/building-index.md)
677
- - [How to reproduce supporting methods?](./docs/reproduce_experiment.md)
682
+ - [How should I set different experimental parameters?](./docs/original_docs/configuration.md)
683
+ - [How to build my own corpus, such as a specific segmented Wikipedia?](./docs/original_docs/process-wiki.md)
684
+ - [How to index my own corpus?](./docs/original_docs/building-index.md)
685
+ - [How to reproduce supporting methods?](./docs/original_docs/reproduce_experiment.md)
678
686
 
679
687
  ## :bookmark: License
680
688
 
@@ -24,7 +24,7 @@
24
24
  </h4>
25
25
 
26
26
 
27
- FlashRAG is a Python toolkit for the reproduction and development of Retrieval Augmented Generation (RAG) research. Our toolkit includes 36 pre-processed benchmark RAG datasets and 16 state-of-the-art RAG algorithms.
27
+ FlashRAG is a Python toolkit for the reproduction and development of Retrieval Augmented Generation (RAG) research. Our toolkit includes 36 pre-processed benchmark RAG datasets and 17 state-of-the-art RAG algorithms.
28
28
 
29
29
  <p align="center">
30
30
  <img src="asset/framework.jpg">
@@ -58,7 +58,7 @@ https://github.com/user-attachments/assets/8ca00873-5df2-48a7-b853-89e7b18bc6e9
58
58
 
59
59
  - **Comprehensive Benchmark Datasets**: A collection of 36 pre-processed RAG benchmark datasets to test and validate RAG models' performances.
60
60
 
61
- - **Pre-implemented Advanced RAG Algorithms**: Features 16 advancing RAG algorithms with reported results, based on our framework. Easily reproducing results under different settings.
61
+ - **Pre-implemented Advanced RAG Algorithms**: Features 17 advancing RAG algorithms with reported results, based on our framework. Easily reproducing results under different settings.
62
62
 
63
63
  - **Efficient Preprocessing Stage**: Simplifies the RAG workflow preparation by providing various scripts like corpus processing for retrieval, retrieval index building, and pre-retrieval of documents.
64
64
 
@@ -79,17 +79,19 @@ FlashRAG is still under development and there are many issues and room for impro
79
79
  - [ ] Enhance code adaptability and readability
80
80
 
81
81
  ## :page_with_curl: Changelog
82
+ [25/03/21] We have added support for **Reasoning Pipeline**, which is a new paradigm that combines reasoning ability and retrieval, representing work that includes [Search-o1](https://github.com/sunnynexus/Search-o1), [R1-Searcher](https://github.com/SsmallSong/R1-Searcher), [ReSearch](https://github.com/Agent-RL/ReSearch). We evaluate the performance of the pipeline on various RAG benchmarks, it can achieve F1 scores close to 60 on multi hop inference datasets such as Hotpotqe. See it in [**result table**](#robot-supporting-methods).
83
+
82
84
  [25/02/24] 🔥🔥🔥 We have added support for **multimodal RAG**, including [**MLLMs like Llava, Qwen, InternVL**](https://ruc-nlpir.github.io/FlashRAG/#/zh-cn/component/generator?id=%e5%a4%9a%e6%a8%a1%e6%80%81%e7%94%9f%e6%88%90%e5%99%a8), and various [**multimodal retrievers with Clip architecture**](https://ruc-nlpir.github.io/FlashRAG/#/zh-cn/component/retriever?id=%e5%a4%9a%e6%a8%a1%e6%80%81%e6%a3%80%e7%b4%a2%e5%99%a8). More information can be found in our new version of arxiv article and our documentation. Try it!
83
85
 
84
86
  [25/01/21] Our technical paper [FlashRAG: A Python Toolkit for Efficient RAG Research](https://arxiv.org/abs/2405.13576) is honored to have been accepted to the Resource Track of the 2025 **ACM Web Conference (WWW 2025)**. Please Check it out!
85
87
 
86
88
  [25/01/12] Introduce <strong>FlashRAG-UI</strong>, an easy to use interface. You can easily and quickly configure and experience the supported RAG methods and evaluate them on the benchmarks.
87
89
 
88
- [25/01/11] We have added support for a new method [<u>RQRAG</u>](https://arxiv.org/abs/2404.00610) method, see it in [**reproduce_experiment**](docs/reproduce_experiment.md).
90
+ [25/01/11] We have added support for a new method [<u>RQRAG</u>](https://arxiv.org/abs/2404.00610) method, see it in [**reproduce_experiment**](docs/original_docs/reproduce_experiment.md).
89
91
 
90
- [25/01/07] We have currently support the aggregation of multiple retrievers, see it in [**multi retriever usage**](https://github.com/RUC-NLPIR/FlashRAG/blob/main/docs/multi_retriever_usage.md).
92
+ [25/01/07] We have currently support the aggregation of multiple retrievers, see it in [**multi retriever usage**](https://github.com/RUC-NLPIR/FlashRAG/blob/main/docs/original_docs/multi_retriever_usage.md).
91
93
 
92
- [25/01/07] We have integrated a very flexible and lightweight corpus chunking library [**Chunkie**](https://github.com/chonkie-ai/chonkie?tab=readme-ov-file#usage), which supports various custom chunking methods (tokens, sentences, semantic, etc.). Use it in [<u>chunking doc corpus</u>](docs/chunk-doc-corpus.md).
94
+ [25/01/07] We have integrated a very flexible and lightweight corpus chunking library [**Chunkie**](https://github.com/chonkie-ai/chonkie?tab=readme-ov-file#usage), which supports various custom chunking methods (tokens, sentences, semantic, etc.). Use it in [<u>chunking doc corpus</u>](docs/original_docs/chunk-doc-corpus.md).
93
95
 
94
96
  [24/10/21] We have released a version based on the Paddle framework that supports Chinese hardware platforms. Please refer to [FlashRAG Paddle](https://github.com/RUC-NLPIR/FlashRAG-Paddle) for details.
95
97
 
@@ -108,7 +110,7 @@ FlashRAG is still under development and there are many issues and room for impro
108
110
 
109
111
  [24/07/17] Due to some unknown issues with HuggingFace, our original dataset link has been invalid. We have updated it. Please check the [new link](https://huggingface.co/datasets/RUC-NLPIR/FlashRAG_datasets/) if you encounter any problems.
110
112
 
111
- [24/07/06] We add support for a new method: [<u>Trace</u>](https://arxiv.org/abs/2406.11460), which refine text by constructing a knowledge graph. See it [<u>results</u>](#robot-supporting-methods) and [<u>details</u>](./docs/baseline_details.md).
113
+ [24/07/06] We add support for a new method: [<u>Trace</u>](https://arxiv.org/abs/2406.11460), which refine text by constructing a knowledge graph. See it [<u>results</u>](#robot-supporting-methods) and [<u>details</u>](./docs/original_docs/baseline_details.md).
112
114
 
113
115
  [24/06/19] We add support for a new method: [<u>IRCoT</u>](https://arxiv.org/abs/2212.10509), and update the [<u>result table</u>](#robot-supporting-methods).
114
116
 
@@ -116,9 +118,9 @@ FlashRAG is still under development and there are many issues and room for impro
116
118
 
117
119
  [24/06/11] We have integrated `sentence transformers` in the retriever module. Now it's easier to use the retriever without setting pooling methods.
118
120
 
119
- [24/06/05] We have provided detailed document for reproducing existing methods (see [how to reproduce](./docs/reproduce_experiment.md), [baseline details](./docs/baseline_details.md)), and [<u>configurations settings</u>](./docs/configuration.md).
121
+ [24/06/05] We have provided detailed document for reproducing existing methods (see [how to reproduce](./docs/original_docs/reproduce_experiment.md), [baseline details](./docs/original_docs/baseline_details.md)), and [<u>configurations settings</u>](./docs/original_docs/configuration.md).
120
122
 
121
- [24/06/02] We have provided an introduction of FlashRAG for beginners, see [<u>an introduction to flashrag</u>](./docs/introduction_for_beginners_en.md) ([<u>中文版</u>](./docs/introduction_for_beginners_zh.md) [<u>한국어</u>](./docs/introduction_for_beginners_kr.md)).
123
+ [24/06/02] We have provided an introduction of FlashRAG for beginners, see [<u>an introduction to flashrag</u>](./docs/original_docs/introduction_for_beginners_en.md) ([<u>中文版</u>](./docs/original_docs/introduction_for_beginners_zh.md) [<u>한국어</u>](./docs/original_docs/introduction_for_beginners_kr.md)).
122
124
 
123
125
  [24/05/31] We supported Openai-series models as generator.
124
126
 
@@ -183,7 +185,7 @@ To build an index, you first need to save your corpus as a `jsonl` file with eac
183
185
  {"id": "1", "contents": "..."}
184
186
  ```
185
187
 
186
- If you want to use Wikipedia as your corpus, you can refer to our documentation [Processing Wikipedia](./docs/process-wiki.md) to convert it into an indexable format.
188
+ If you want to use Wikipedia as your corpus, you can refer to our documentation [Processing Wikipedia](./docs/original_docs/process-wiki.md) to convert it into an indexable format.
187
189
 
188
190
  ### Index Construction
189
191
 
@@ -272,7 +274,7 @@ my_config = Config(
272
274
  config_dict = config_dict
273
275
  ```
274
276
 
275
- We provide comprehensive guidance on how to set configurations, you can see our [<u>configuration guidance</u>](./docs/configuration.md).
277
+ We provide comprehensive guidance on how to set configurations, you can see our [<u>configuration guidance</u>](./docs/original_docs/configuration.md).
276
278
  You can also refer to the [<u>basic yaml file</u>](./flashrag/config/basic_config.yaml) we provide to set your own parameters.
277
279
 
278
280
  Next, load the corresponding dataset and initialize the pipeline. The components in the pipeline will be automatically loaded.
@@ -343,11 +345,11 @@ class ToyPipeline(BasicPipeline):
343
345
  return dataset
344
346
  ```
345
347
 
346
- Please first understand the input and output forms of the components you need to use from our [<u>documentation</u>](./docs/basic_usage.md).
348
+ Please first understand the input and output forms of the components you need to use from our [<u>documentation</u>](./docs/original_docs/basic_usage.md).
347
349
 
348
350
  ### Just use components
349
351
 
350
- If you already have your own code and only want to use our components to embed the original code, you can refer to the [<u>basic introduction of the components</u>](./docs/basic_usage.md) to obtain the input and output formats of each component.
352
+ If you already have your own code and only want to use our components to embed the original code, you can refer to the [<u>basic introduction of the components</u>](./docs/original_docs/basic_usage.md) to obtain the input and output formats of each component.
351
353
 
352
354
  ## :gear: Components
353
355
 
@@ -465,7 +467,7 @@ In each category, we have implemented corresponding common pipelines. Some pipel
465
467
  <td>Ranking and merging generated results based on each document</td>
466
468
  </tr>
467
469
  <tr>
468
- <td rowspan="5">Loop</td>
470
+ <td rowspan="6">Loop</td>
469
471
  <td>Iterative Pipeline</td>
470
472
  <td>Alternating retrieval and generation</td>
471
473
  </tr>
@@ -485,6 +487,10 @@ In each category, we have implemented corresponding common pipelines. Some pipel
485
487
  <td>IRCoT Pipeline</td>
486
488
  <td>Integrate retrieval process with CoT</td>
487
489
  </tr>
490
+ <tr>
491
+ <td>Reasoning Pipeline</td>
492
+ <td>Reasoning with retrieval</td>
493
+ </tr>
488
494
  </tbody>
489
495
  </table>
490
496
 
@@ -537,11 +543,11 @@ We have implemented 15 works with a consistent setting of:
537
543
 
538
544
  - **Generator:** LLAMA3-8B-instruct with input length of 2048
539
545
  - **Retriever:** e5-base-v2 as embedding model, retrieve 5 docs per query
540
- - **Prompt:** A consistent default prompt, template can be found in the [<u>method details</u>](./docs/baseline_details.md).
546
+ - **Prompt:** A consistent default prompt, template can be found in the [<u>method details</u>](./docs/original_docs/baseline_details.md).
541
547
 
542
548
  For open-source methods, we implemented their processes using our framework. For methods where the author did not provide source code, we will try our best to follow the methods in the original paper for implementation.
543
549
 
544
- For necessary settings and hyperparameters specific to some methods, we have documented them in the **specific settings** column. For more details, please consult our [<u>reproduce guidance</u>](./docs/reproduce_experiment.md) and [<u>method details</u>](./docs/baseline_details.md).
550
+ For necessary settings and hyperparameters specific to some methods, we have documented them in the **specific settings** column. For more details, please consult our [<u>reproduce guidance</u>](./docs/original_docs/reproduce_experiment.md) and [<u>method details</u>](./docs/original_docs/baseline_details.md).
545
551
 
546
552
  It’s important to note that, to ensure consistency, we have utilized a uniform setting. However, this setting may differ from the original setting of the method, leading to variations in results compared to the original outcomes.
547
553
 
@@ -565,6 +571,8 @@ It’s important to note that, to ensure consistency, we have utilized a uniform
565
571
  | [Iter-Retgen](https://arxiv.org/abs/2305.15294), [ITRG](https://arxiv.org/abs/2310.05149) | Loop | 36.8 | 60.1 | 38.3 | 21.6 | 37.9 | 18.2 | |
566
572
  | [IRCoT](https://aclanthology.org/2023.acl-long.557.pdf) | Loop | 33.3 | 56.9 | 41.5 | 32.4 | 45.6 | 20.7 | |
567
573
  | [RQRAG](https://arxiv.org/abs/2404.00610) | Loop | 32.6 | 52.5 | 33.5 | 35.8 | 46.4 | 26.2 | Use trained rqrag-llama2-7B |
574
+ | [R1-Searcher](https://arxiv.org/pdf/2503.05592) | Reasoning | 37.3 | 58.7 | 59.5 | 55.5 | 43.9 | 27.5 | Using trained Qwen2.5-7B model |
575
+
568
576
 
569
577
 
570
578
  ## :notebook: Supporting Datasets & Document Corpus
@@ -636,7 +644,7 @@ Our toolkit supports jsonl format for retrieval document collections, with the f
636
644
 
637
645
  The `contents` key is essential for building the index. For documents that include both text and title, we recommend setting the value of `contents` to `{title}\n{text}`. The corpus file can also contain other keys to record additional characteristics of the documents.
638
646
 
639
- In the academic research, Wikipedia and MS MARCO are the most commonly used retrieval document collections. For Wikipedia, we provide a [<u>comprehensive script</u>](./docs/process-wiki.md) to process any Wikipedia dump into a clean corpus. Additionally, various processed versions of the Wikipedia corpus are available in many works, and we have listed some reference links.
647
+ In the academic research, Wikipedia and MS MARCO are the most commonly used retrieval document collections. For Wikipedia, we provide a [<u>comprehensive script</u>](./docs/original_docs/process-wiki.md) to process any Wikipedia dump into a clean corpus. Additionally, various processed versions of the Wikipedia corpus are available in many works, and we have listed some reference links.
640
648
 
641
649
  For MS MARCO, it is already processed upon release and can be directly downloaded from its [<u>hosting link</u>](https://huggingface.co/datasets/Tevatron/msmarco-passage-corpus) on Hugging Face.
642
650
 
@@ -654,10 +662,10 @@ The index was created using the e5-base-v2 retriever on our uploaded wiki18_100w
654
662
 
655
663
  ## :raised_hands: Additional FAQs
656
664
 
657
- - [How should I set different experimental parameters?](./docs/configuration.md)
658
- - [How to build my own corpus, such as a specific segmented Wikipedia?](./docs/process-wiki.md)
659
- - [How to index my own corpus?](./docs/building-index.md)
660
- - [How to reproduce supporting methods?](./docs/reproduce_experiment.md)
665
+ - [How should I set different experimental parameters?](./docs/original_docs/configuration.md)
666
+ - [How to build my own corpus, such as a specific segmented Wikipedia?](./docs/original_docs/process-wiki.md)
667
+ - [How to index my own corpus?](./docs/original_docs/building-index.md)
668
+ - [How to reproduce supporting methods?](./docs/original_docs/reproduce_experiment.md)
661
669
 
662
670
  ## :bookmark: License
663
671
 
@@ -66,6 +66,8 @@ retrieval_cache_path: ~ # path to the retrieval cache
66
66
  retrieval_pooling_method: ~ # set automatically if not provided
67
67
  bm25_backend: bm25s # pyserini, bm25s
68
68
  use_sentence_transformer: False
69
+ silent_retrieval: True # whether to silent the retrieval process
70
+
69
71
 
70
72
  use_reranker: False # whether to use reranker
71
73
  rerank_model_name: ~ # same as retrieval_method
@@ -48,6 +48,13 @@ class Item:
48
48
  except AttributeError:
49
49
  raise AttributeError(f"Attribute `{attr_name}` not found")
50
50
 
51
+ def __setattr__(self, attr_name: str, value: Any) -> None:
52
+ predefined_attrs = ["id", "question", "golden_answers", "metadata", "output", "choices", 'data']
53
+ if attr_name in predefined_attrs:
54
+ super().__setattr__(attr_name, value)
55
+ else:
56
+ self.update_output(attr_name, value)
57
+
51
58
  def to_dict(self) -> Dict[str, Any]:
52
59
  """Convert all information within the data sample into a dict. Information generated
53
60
  during the inference will be saved into output field.
@@ -35,7 +35,7 @@ class OpenaiGenerator:
35
35
  try:
36
36
  self.tokenizer = tiktoken.encoding_for_model(self.model_name)
37
37
  except Exception as e:
38
- print("Error: ", e)
38
+ print("Warning: ", e)
39
39
  warnings.warn("This model is not supported by tiktoken. Use gpt-3.5-turbo instead.")
40
40
  self.tokenizer = tiktoken.encoding_for_model('gpt-3.5-turbo')
41
41
  @property
@@ -1,4 +1,5 @@
1
1
  from flashrag.pipeline.mm_pipeline import *
2
2
  from flashrag.pipeline.pipeline import *
3
3
  from flashrag.pipeline.branching_pipeline import REPLUGPipeline, SuRePipeline
4
- from flashrag.pipeline.active_pipeline import IterativePipeline, SelfRAGPipeline, FLAREPipeline, SelfAskPipeline, IRCOTPipeline, RQRAGPipeline
4
+ from flashrag.pipeline.active_pipeline import IterativePipeline, SelfRAGPipeline, FLAREPipeline, SelfAskPipeline, IRCOTPipeline, RQRAGPipeline
5
+ from flashrag.pipeline.reasoning_pipeline import ReasoningPipeline
@@ -0,0 +1,154 @@
1
+ import re
2
+ from tqdm import tqdm
3
+ from typing import List, Tuple, Dict
4
+ import math
5
+ import numpy as np
6
+ import copy
7
+ from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
8
+ from flashrag.utils import get_retriever, get_generator, selfask_pred_parse, ircot_pred_parse
9
+ from flashrag.pipeline import BasicPipeline
10
+ from flashrag.dataset.utils import get_batch_dataset, merge_batch_dataset
11
+ from flashrag.prompt import PromptTemplate
12
+ from flashrag.utils.utils import extract_between
13
+
14
+ class ReasoningPipeline(BasicPipeline):
15
+ system_prompt=""
16
+ user_prompt=(
17
+ "The User asks a question, and the Assistant solves it.\n"
18
+ "The Assistant first thinks about the reasoning process in the mind and then provides the User with the final answer.\n"
19
+ "The output format of reasoning process and final answer are enclosed within <think> </think> and <answer> </answer> tags, "
20
+ "respectively, i.e., \"<think> reasoning process here </think>\\n\\n<answer> final answer here </answer>\".\n"
21
+ "During the thinking process, the Assistant can perform searching for uncertain knowledge if necessary with the format of "
22
+ "\"<|begin_of_query|> search query (only keywords) here <|end_of_query|>\". **A query must involve only a single triple**.\n"
23
+ "Then, the system will provide the Assistant with helpful information with the format of "
24
+ "\"<|begin_of_documents|> ...search results... <|end_of_documents|>\".\n\n"
25
+ "User:{question}\n"
26
+ "Assistant: <think>"
27
+ )
28
+
29
+
30
+ def __init__(self, config,
31
+ prompt_template=None,
32
+ max_retrieval_num=5,
33
+ begin_of_query_token="<|begin_of_query|>",
34
+ end_of_query_token="<|end_of_query|>",
35
+ begin_of_documents_token="<|begin_of_documents|>",
36
+ end_of_documents_token="<|end_of_documents|>",
37
+ begin_of_answer_token="<answer>",
38
+ end_of_answer_token='</answer>',
39
+ retriever=None,
40
+ generator=None
41
+ ):
42
+ if prompt_template is None:
43
+ prompt_template = PromptTemplate(
44
+ config=config,
45
+ system_prompt=self.system_prompt,
46
+ user_prompt=self.user_prompt
47
+ )
48
+ super().__init__(config, prompt_template)
49
+
50
+ if generator is None:
51
+ self.generator = get_generator(config)
52
+ else:
53
+ self.generator = generator
54
+ if retriever is None:
55
+ self.retriever = get_retriever(config)
56
+ else:
57
+ self.retriever = retriever
58
+
59
+ self.max_retrieval_num = max_retrieval_num
60
+
61
+ self.begin_of_query_token = begin_of_query_token
62
+ self.end_of_query_token = end_of_query_token
63
+ self.begin_of_documents_token = begin_of_documents_token
64
+ self.end_of_documents_token = end_of_documents_token
65
+ self.begin_of_answer_token = begin_of_answer_token
66
+ self.end_of_answer_token = end_of_answer_token
67
+
68
+ self.stop_tokens = ["<|im_end|>", "<|endoftext|>", self.end_of_answer_token, self.end_of_query_token]
69
+
70
+ def _retrieved_docs_to_string(self, retrieved_docs: List[Dict]):
71
+ format_doc_string = ""
72
+ for idx, doc in enumerate(retrieved_docs):
73
+ contents = doc['contents']
74
+ title = contents.split('\n')[0]
75
+ text = '\n'.join(contents.split('\n')[1:])
76
+ doc_string = f"Title: {title} Text: {text}"
77
+ doc_string = re.sub(r'^\d+\s+', '', doc_string)
78
+ format_doc_string += f'({idx+1}){doc_string}\n'
79
+ format_doc_string = f'\n\n{self.begin_of_documents_token}\n{format_doc_string}\n{self.end_of_documents_token}\n\n'
80
+ return format_doc_string
81
+
82
+ def run(self, dataset, do_eval=True, pred_process_fun=None):
83
+ prompts = [self.prompt_template.get_string(question=question) for question in dataset.question]
84
+ dataset.update_output('prompt', prompts)
85
+ dataset.update_output('finish_flag', [False] * len(prompts))
86
+ dataset.update_output('retrieval_results', [{}] * len(prompts))
87
+ dataset.update_output('retrieved_times', [0] * len(prompts))
88
+
89
+ # Logic of reasoning
90
+ for current_step_idx in range(self.max_retrieval_num + 1):
91
+ exist_items = [item for item in dataset if item.finish_flag == False]
92
+ exist_prompts = [item.prompt for item in exist_items]
93
+
94
+ print(f"Current step: {current_step_idx}, exist_items: {len(exist_items)}")
95
+
96
+ if len(exist_items) == 0:
97
+ print("All prompts are finished")
98
+ break
99
+ if current_step_idx == self.max_retrieval_num:
100
+ print("Max retrieval number reached")
101
+ for item in exist_items:
102
+ item.pred = 'No valid answer found'
103
+ item.finish_flag = True
104
+ item.finish_reason = 'Reach max retrieval number'
105
+ break
106
+
107
+ step_outputs = self.generator.generate(exist_prompts, stop=self.stop_tokens)
108
+ step_query_list = [] # store generated queries for retrieval
109
+
110
+ # parse each sample's step output
111
+ for item, step_output in zip(exist_items, step_outputs):
112
+ item.prompt = item.prompt + step_output.strip()
113
+ if self.end_of_answer_token in step_output and step_output.endswith(self.end_of_answer_token):
114
+ item.pred = str(extract_between(step_output, self.begin_of_answer_token, self.end_of_answer_token))
115
+ item.finish_flag = True
116
+ item.finish_reason = "Finished"
117
+
118
+ elif self.begin_of_query_token in step_output and step_output.endswith(self.end_of_query_token):
119
+ query = extract_between(step_output, self.begin_of_query_token, self.end_of_query_token)
120
+ if query is not None:
121
+ step_query_list.append({'item': item, 'query': query})
122
+ else:
123
+ item.pred = 'No valid answer found'
124
+ item.finish_flag = True
125
+ item.finish_reason = 'Query instruction error'
126
+
127
+ else:
128
+ item.pred = step_output.strip()
129
+ item.finish_flag = True
130
+ item.finish_reason = 'Normal finish without answer pattern'
131
+
132
+ # do retrieval and add retrieved docs to prompt
133
+ if len(step_query_list) > 0:
134
+ retrieved_docs = self.retriever.batch_search([it['query'] for it in step_query_list])
135
+ for it, item_retrieved_docs in zip(step_query_list, retrieved_docs):
136
+ item = it['item']
137
+ query = it['query']
138
+ item.retrieval_results[item.retrieved_times] = {'query': query,'docs': copy.copy(item_retrieved_docs)}
139
+ #item.retrieved_docs += [item_retrieved_docs]
140
+ format_doc_string = self._retrieved_docs_to_string(item_retrieved_docs)
141
+ item.prompt += format_doc_string
142
+ item.retrieved_times += 1
143
+
144
+ dataset = self.evaluate(dataset, do_eval=do_eval, pred_process_fun=pred_process_fun)
145
+ return dataset
146
+
147
+
148
+
149
+
150
+
151
+
152
+
153
+
154
+
@@ -24,13 +24,14 @@ class Encoder:
24
24
  Encodes a list of queries into embeddings.
25
25
  """
26
26
 
27
- def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16, instruction):
27
+ def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16, instruction, silent=False):
28
28
  self.model_name = model_name
29
29
  self.model_path = model_path
30
30
  self.pooling_method = pooling_method
31
31
  self.max_length = max_length
32
32
  self.use_fp16 = use_fp16
33
33
  self.instruction = instruction
34
+ self.silent = silent
34
35
  self.gpu_num = torch.cuda.device_count()
35
36
  self.model, self.tokenizer = load_model(model_path=model_path, use_fp16=use_fp16)
36
37
 
@@ -67,7 +68,7 @@ class Encoder:
67
68
  @torch.inference_mode()
68
69
  def encode(self, query_list: List[str], batch_size=64, is_query=True) -> np.ndarray:
69
70
  query_emb = []
70
- for i in tqdm(range(0, len(query_list), batch_size), desc="Encoding process: "):
71
+ for i in tqdm(range(0, len(query_list), batch_size), desc="Encoding process: ", disable=self.silent):
71
72
  query_emb.append(self.single_batch_encode(query_list[i : i + batch_size], is_query))
72
73
  query_emb = np.concatenate(query_emb, axis=0)
73
74
  return query_emb
@@ -98,7 +99,7 @@ class STEncoder:
98
99
  Encodes a list of queries into embeddings using multiple GPUs.
99
100
  """
100
101
 
101
- def __init__(self, model_name, model_path, max_length, use_fp16, instruction):
102
+ def __init__(self, model_name, model_path, max_length, use_fp16, instruction, silent=False):
102
103
  import torch
103
104
  from sentence_transformers import SentenceTransformer
104
105
 
@@ -107,6 +108,7 @@ class STEncoder:
107
108
  self.max_length = max_length
108
109
  self.use_fp16 = use_fp16
109
110
  self.instruction = instruction
111
+ self.silent = silent
110
112
  self.model = SentenceTransformer(
111
113
  model_path, trust_remote_code=True, model_kwargs={"torch_dtype": torch.float16 if use_fp16 else torch.float}
112
114
  )
@@ -115,7 +117,7 @@ class STEncoder:
115
117
  def encode(self, query_list: Union[List[str], str], batch_size=64, is_query=True) -> np.ndarray:
116
118
  query_list = parse_query(self.model_name, query_list, self.instruction, is_query)
117
119
  query_emb = self.model.encode(
118
- query_list, batch_size=batch_size, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=True
120
+ query_list, batch_size=batch_size, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=not self.silent
119
121
  )
120
122
  query_emb = query_emb.astype(np.float32, order="C")
121
123
 
@@ -130,7 +132,7 @@ class STEncoder:
130
132
  pool,
131
133
  normalize_embeddings=True,
132
134
  batch_size=batch_size,
133
- show_progress_bar=True,
135
+ show_progress_bar=not self.silent,
134
136
  )
135
137
  self.model.stop_multi_process_pool(pool)
136
138
  query_emb = query_emb.astype(np.float32, order="C")
@@ -190,7 +192,7 @@ class ClipEncoder:
190
192
  if not isinstance(query_list, list):
191
193
  query_list = [query_list]
192
194
  query_emb = []
193
- for i in tqdm(range(0, len(query_list), batch_size), desc="Encoding process: "):
195
+ for i in tqdm(range(0, len(query_list), batch_size), desc="Encoding process: ", disable=self.silent):
194
196
  query_emb.append(self.single_batch_encode(query_list[i : i + batch_size], modal))
195
197
  query_emb = np.concatenate(query_emb, axis=0)
196
198
  return query_emb