flashrag-dev 0.1.4.dev20250319__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.dev20250319 → flashrag_dev-0.1.4.dev20250321}/PKG-INFO +12 -4
  2. flashrag_dev-0.1.4.dev20250319/flashrag_dev.egg-info/PKG-INFO → flashrag_dev-0.1.4.dev20250321/README.md +11 -20
  3. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/config/basic_config.yaml +2 -0
  4. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/dataset/dataset.py +7 -0
  5. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/pipeline/__init__.py +2 -1
  6. flashrag_dev-0.1.4.dev20250321/flashrag/pipeline/reasoning_pipeline.py +154 -0
  7. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/retriever/encoder.py +8 -6
  8. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/retriever/retriever.py +97 -63
  9. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/retriever/utils.py +11 -4
  10. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/utils/utils.py +8 -0
  11. flashrag_dev-0.1.4.dev20250321/flashrag/version.py +1 -0
  12. flashrag_dev-0.1.4.dev20250319/README.md → flashrag_dev-0.1.4.dev20250321/flashrag_dev.egg-info/PKG-INFO +28 -3
  13. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag_dev.egg-info/SOURCES.txt +1 -0
  14. flashrag_dev-0.1.4.dev20250319/flashrag/version.py +0 -1
  15. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/LICENSE +0 -0
  16. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/__init__.py +0 -0
  17. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/config/__init__.py +0 -0
  18. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/config/config.py +0 -0
  19. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/dataset/__init__.py +0 -0
  20. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/dataset/utils.py +0 -0
  21. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/evaluator/__init__.py +0 -0
  22. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/evaluator/_bleu.py +0 -0
  23. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/evaluator/evaluator.py +0 -0
  24. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/evaluator/metrics.py +0 -0
  25. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/evaluator/utils.py +0 -0
  26. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/generator/__init__.py +0 -0
  27. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/generator/fid.py +0 -0
  28. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/generator/generator.py +0 -0
  29. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/generator/multimodal_generator.py +0 -0
  30. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/generator/openai_generator.py +0 -0
  31. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/generator/stop_word_criteria.py +0 -0
  32. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/generator/utils.py +0 -0
  33. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/judger/__init__.py +0 -0
  34. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/judger/judger.py +0 -0
  35. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/pipeline/active_pipeline.py +0 -0
  36. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/pipeline/branching_pipeline.py +0 -0
  37. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/pipeline/mm_pipeline.py +0 -0
  38. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/pipeline/pipeline.py +0 -0
  39. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/pipeline/replug_utils.py +0 -0
  40. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/prompt/__init__.py +0 -0
  41. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/prompt/base_prompt.py +0 -0
  42. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/prompt/mm_prompt.py +0 -0
  43. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/prompt/selfask_examplars.py +0 -0
  44. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/prompt/trace_examplars.py +0 -0
  45. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/refiner/__init__.py +0 -0
  46. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/refiner/kg_refiner.py +0 -0
  47. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/refiner/llmlingua_compressor.py +0 -0
  48. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/refiner/refiner.py +0 -0
  49. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/refiner/selective_context_compressor.py +0 -0
  50. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/retriever/__init__.py +0 -0
  51. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/retriever/__main__.py +0 -0
  52. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/retriever/index_builder.py +0 -0
  53. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/retriever/reranker.py +0 -0
  54. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/utils/__init__.py +0 -0
  55. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/utils/constants.py +0 -0
  56. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag/utils/pred_parse.py +0 -0
  57. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag_dev.egg-info/dependency_links.txt +0 -0
  58. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag_dev.egg-info/requires.txt +0 -0
  59. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/flashrag_dev.egg-info/top_level.txt +0 -0
  60. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/pyproject.toml +0 -0
  61. {flashrag_dev-0.1.4.dev20250319 → flashrag_dev-0.1.4.dev20250321}/setup.cfg +0 -0
  62. {flashrag_dev-0.1.4.dev20250319 → 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.dev20250319
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,6 +96,8 @@ 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!
@@ -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
 
@@ -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
@@ -1,20 +1,3 @@
1
- Metadata-Version: 2.1
2
- Name: flashrag-dev
3
- Version: 0.1.4.dev20250319
4
- Summary: A library for efficient Retrieval-Augmented Generation research
5
- Home-page: https://github.com/RUC-NLPIR/FlashRAG
6
- Author: Jiajie Jin, Yutao Zhu, Chenghao Zhang, Xinyu Yang, Zhicheng Dou
7
- Author-email: jinjiajie@ruc.edu.cn
8
- License: MIT License
9
- Requires-Python: >=3.9
10
- Description-Content-Type: text/markdown
11
- Provides-Extra: core
12
- Provides-Extra: retriever
13
- Provides-Extra: generator
14
- Provides-Extra: multimodal
15
- Provides-Extra: full
16
- License-File: LICENSE
17
-
18
1
  # <div align="center">⚡FlashRAG: A Python Toolkit for Efficient RAG Research<div>
19
2
  \[ English | [中文](README_zh.md) \]
20
3
  <div align="center">
@@ -41,7 +24,7 @@ License-File: LICENSE
41
24
  </h4>
42
25
 
43
26
 
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.
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.
45
28
 
46
29
  <p align="center">
47
30
  <img src="asset/framework.jpg">
@@ -75,7 +58,7 @@ https://github.com/user-attachments/assets/8ca00873-5df2-48a7-b853-89e7b18bc6e9
75
58
 
76
59
  - **Comprehensive Benchmark Datasets**: A collection of 36 pre-processed RAG benchmark datasets to test and validate RAG models' performances.
77
60
 
78
- - **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.
79
62
 
80
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.
81
64
 
@@ -96,6 +79,8 @@ FlashRAG is still under development and there are many issues and room for impro
96
79
  - [ ] Enhance code adaptability and readability
97
80
 
98
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
+
99
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!
100
85
 
101
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!
@@ -482,7 +467,7 @@ In each category, we have implemented corresponding common pipelines. Some pipel
482
467
  <td>Ranking and merging generated results based on each document</td>
483
468
  </tr>
484
469
  <tr>
485
- <td rowspan="5">Loop</td>
470
+ <td rowspan="6">Loop</td>
486
471
  <td>Iterative Pipeline</td>
487
472
  <td>Alternating retrieval and generation</td>
488
473
  </tr>
@@ -502,6 +487,10 @@ In each category, we have implemented corresponding common pipelines. Some pipel
502
487
  <td>IRCoT Pipeline</td>
503
488
  <td>Integrate retrieval process with CoT</td>
504
489
  </tr>
490
+ <tr>
491
+ <td>Reasoning Pipeline</td>
492
+ <td>Reasoning with retrieval</td>
493
+ </tr>
505
494
  </tbody>
506
495
  </table>
507
496
 
@@ -582,6 +571,8 @@ It’s important to note that, to ensure consistency, we have utilized a uniform
582
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 | |
583
572
  | [IRCoT](https://aclanthology.org/2023.acl-long.557.pdf) | Loop | 33.3 | 56.9 | 41.5 | 32.4 | 45.6 | 20.7 | |
584
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
+
585
576
 
586
577
 
587
578
  ## :notebook: Supporting Datasets & Document Corpus
@@ -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.
@@ -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
@@ -1,6 +1,7 @@
1
1
  import json
2
2
  import os
3
3
  import time
4
+
4
5
  os.environ["TOKENIZERS_PARALLELISM"] = "false"
5
6
  import warnings
6
7
  from typing import List, Dict, Union
@@ -116,7 +117,7 @@ class BaseRetriever:
116
117
  def __init__(self, config):
117
118
  self._config = config
118
119
  self.update_config()
119
-
120
+
120
121
  @property
121
122
  def config(self):
122
123
  return self._config
@@ -125,11 +126,11 @@ class BaseRetriever:
125
126
  def config(self, config_data):
126
127
  self._config = config_data
127
128
  self.update_config()
128
-
129
+
129
130
  def update_config(self):
130
131
  self.update_base_setting()
131
132
  self.update_additional_setting()
132
-
133
+
133
134
  def update_base_setting(self):
134
135
  self.retrieval_method = self._config["retrieval_method"]
135
136
  self.topk = self._config["retrieval_topk"]
@@ -154,8 +155,10 @@ class BaseRetriever:
154
155
  assert self.cache_path is not None
155
156
  with open(self.cache_path, "r") as f:
156
157
  self.cache = json.load(f)
158
+ self.silent = self._config["silent_retrieval"] if "silent_retrieval" in self._config else False
157
159
  def update_additional_setting(self):
158
160
  pass
161
+
159
162
  def _save_cache(self):
160
163
  self.cache = convert_numpy(self.cache)
161
164
 
@@ -163,6 +166,7 @@ class BaseRetriever:
163
166
  if isinstance(obj, np.float32):
164
167
  return float(obj)
165
168
  raise TypeError(f"Type {type(obj)} not serializable")
169
+
166
170
  with open(self.cache_save_path, "w") as f:
167
171
  json.dump(self.cache, f, indent=4, default=custom_serializer)
168
172
 
@@ -220,9 +224,10 @@ class BM25Retriever(BaseTextRetriever):
220
224
  def __init__(self, config, corpus=None):
221
225
  super().__init__(config)
222
226
  self.load_model_corpus(corpus)
227
+
223
228
  def update_additional_setting(self):
224
229
  self.backend = self._config["bm25_backend"]
225
-
230
+
226
231
  def load_model_corpus(self, corpus):
227
232
  if self.backend == "pyserini":
228
233
  # Warning: the method based on pyserini will be deprecated
@@ -237,24 +242,24 @@ class BM25Retriever(BaseTextRetriever):
237
242
  self.corpus = corpus
238
243
  self.max_process_num = 8
239
244
 
240
- is_zh = judge_zh(self.corpus[0]['contents'])
245
+ is_zh = judge_zh(self.corpus[0]["contents"])
241
246
  if is_zh:
242
- self.searcher.set_language('zh')
247
+ self.searcher.set_language("zh")
243
248
  elif self.backend == "bm25s":
244
249
  import Stemmer
245
250
  import bm25s
246
251
 
247
252
  self.corpus = load_corpus(self.corpus_path)
248
- is_zh = judge_zh(self.corpus[0]['contents'])
253
+ is_zh = judge_zh(self.corpus[0]["contents"])
249
254
 
250
255
  self.searcher = bm25s.BM25.load(self.index_path, mmap=True, load_corpus=False)
251
256
  if is_zh:
252
- self.tokenizer = bm25s.tokenization.Tokenizer(stopwords='zh')
257
+ self.tokenizer = bm25s.tokenization.Tokenizer(stopwords="zh")
253
258
  self.tokenizer.load_stopwords(self.index_path)
254
259
  self.tokenizer.load_vocab(self.index_path)
255
260
  else:
256
261
  stemmer = Stemmer.Stemmer("english")
257
- self.tokenizer = bm25s.tokenization.Tokenizer(stopwords='en', stemmer=stemmer)
262
+ self.tokenizer = bm25s.tokenization.Tokenizer(stopwords="en", stemmer=stemmer)
258
263
  self.tokenizer.load_stopwords(self.index_path)
259
264
  self.tokenizer.load_vocab(self.index_path)
260
265
 
@@ -298,7 +303,7 @@ class BM25Retriever(BaseTextRetriever):
298
303
  else:
299
304
  results = load_docs(self.corpus, [hit.docid for hit in hits])
300
305
  elif self.backend == "bm25s":
301
- query_tokens = self.tokenizer.tokenize([query], return_as='tuple', update_vocab=False)
306
+ query_tokens = self.tokenizer.tokenize([query], return_as="tuple", update_vocab=False)
302
307
  results, scores = self.searcher.retrieve(query_tokens, k=num)
303
308
  results = list(results[0])
304
309
  scores = list(scores[0])
@@ -320,12 +325,12 @@ class BM25Retriever(BaseTextRetriever):
320
325
  results.append(item_result)
321
326
  scores.append(item_score)
322
327
  elif self.backend == "bm25s":
323
- query_tokens = self.tokenizer.tokenize(query, return_as='tuple', update_vocab=False)
328
+ query_tokens = self.tokenizer.tokenize(query, return_as="tuple", update_vocab=False)
324
329
  results, scores = self.searcher.retrieve(query_tokens, k=num)
325
330
  else:
326
331
  assert False, "Invalid bm25 backend!"
327
- results = results.tolist() if isinstance(results,np.ndarray) else results
328
- scores = scores.tolist() if isinstance(scores,np.ndarray) else scores
332
+ results = results.tolist() if isinstance(results, np.ndarray) else results
333
+ scores = scores.tolist() if isinstance(scores, np.ndarray) else scores
329
334
  if return_score:
330
335
  return results, scores
331
336
  else:
@@ -337,17 +342,17 @@ class DenseRetriever(BaseTextRetriever):
337
342
 
338
343
  def __init__(self, config: dict, corpus=None):
339
344
  super().__init__(config)
340
-
345
+
341
346
  self.load_corpus(corpus)
342
347
  self.load_index()
343
348
  self.load_model()
344
-
349
+
345
350
  def load_corpus(self, corpus):
346
351
  if corpus is None:
347
352
  self.corpus = load_corpus(self.corpus_path)
348
353
  else:
349
354
  self.corpus = corpus
350
-
355
+
351
356
  def load_index(self):
352
357
  if self.index_path is None or not os.path.exists(self.index_path):
353
358
  raise Warning(f"Index file {self.index_path} does not exist!")
@@ -358,17 +363,16 @@ class DenseRetriever(BaseTextRetriever):
358
363
  co.shard = True
359
364
  self.index = faiss.index_cpu_to_all_gpus(self.index, co=co)
360
365
 
361
-
362
366
  def update_additional_setting(self):
363
367
  self.query_max_length = self._config["retrieval_query_max_length"]
364
- self.pooling_method = self._config['retrieval_pooling_method']
365
- self.use_fp16 = self._config['retrieval_use_fp16']
368
+ self.pooling_method = self._config["retrieval_pooling_method"]
369
+ self.use_fp16 = self._config["retrieval_use_fp16"]
366
370
  self.batch_size = self._config["retrieval_batch_size"]
367
371
  self.instruction = self._config["instruction"]
368
372
 
369
- self.retreival_model_path = self._config['retrieval_model_path']
373
+ self.retreival_model_path = self._config["retrieval_model_path"]
370
374
  self.use_st = self._config["use_sentence_transformer"]
371
- self.use_faiss_gpu = self._config['faiss_gpu']
375
+ self.use_faiss_gpu = self._config["faiss_gpu"]
372
376
 
373
377
  def load_model(self):
374
378
  if self.use_st:
@@ -378,8 +382,11 @@ class DenseRetriever(BaseTextRetriever):
378
382
  max_length = self.query_max_length,
379
383
  use_fp16 = self.use_fp16,
380
384
  instruction = self.instruction,
385
+ silent = self.silent
381
386
  )
382
387
  else:
388
+ # check pooling method
389
+ self._check_pooling_method(self.retrieval_model_path, self.pooling_method)
383
390
  self.encoder = Encoder(
384
391
  model_name = self.retrieval_method,
385
392
  model_path = self.retreival_model_path,
@@ -387,6 +394,7 @@ class DenseRetriever(BaseTextRetriever):
387
394
  max_length = self.query_max_length,
388
395
  use_fp16 = self.use_fp16,
389
396
  instruction = self.instruction,
397
+ silent = self.silent
390
398
  )
391
399
 
392
400
  def _search(self, query: str, num: int = None, return_score=False):
@@ -413,11 +421,8 @@ class DenseRetriever(BaseTextRetriever):
413
421
 
414
422
  results = []
415
423
  scores = []
416
-
417
424
  emb = self.encoder.encode(query, batch_size=batch_size, is_query=True)
418
- print("Begin faiss searching...")
419
425
  scores, idxs = self.index.search(emb, k=num)
420
- print("End faiss searching")
421
426
  scores = scores.tolist()
422
427
  idxs = idxs.tolist()
423
428
 
@@ -459,6 +464,7 @@ class MultiModalRetriever(BaseRetriever):
459
464
  self.encoder = ClipEncoder(
460
465
  model_name=self.retrieval_method,
461
466
  model_path=config["retrieval_model_path"],
467
+ silent=self.silent
462
468
  )
463
469
 
464
470
  def _judge_input_modal(self, query):
@@ -480,10 +486,12 @@ class MultiModalRetriever(BaseRetriever):
480
486
  )
481
487
  if query_modal == "image" and isinstance(query, str):
482
488
  from PIL import Image
489
+
483
490
  if os.path.exists(query):
484
491
  query = Image.open(query)
485
492
  else:
486
493
  import requests
494
+
487
495
  query = Image.open(requests.get(query, stream=True).raw)
488
496
 
489
497
  query_emb = self.encoder.encode(query, modal=query_modal)
@@ -511,6 +519,7 @@ class MultiModalRetriever(BaseRetriever):
511
519
  if query_modal == "image" and isinstance(query[0], str):
512
520
  from PIL import Image
513
521
  import requests
522
+
514
523
  if os.path.exists(query[0]):
515
524
  query = [Image.open(q) for q in query]
516
525
  else:
@@ -519,7 +528,7 @@ class MultiModalRetriever(BaseRetriever):
519
528
  results = []
520
529
  scores = []
521
530
 
522
- for start_idx in tqdm(range(0, len(query), batch_size), desc="Retrieval process: "):
531
+ for start_idx in tqdm(range(0, len(query), batch_size), desc="Retrieval process: ", disable=self.silent):
523
532
  query_batch = query[start_idx : start_idx + batch_size]
524
533
  batch_emb = self.encoder.encode(query_batch, modal=query_modal)
525
534
  batch_scores, batch_idxs = self.index_dict[target_modal].search(batch_emb, k=num)
@@ -547,10 +556,10 @@ class MultiRetrieverRouter:
547
556
  self.retriever_list = self.load_all_retriever(config)
548
557
  self.config = config
549
558
 
550
- if self.merge_method == 'rerank':
551
- config['multi_retriever_setting']['rerank_topk'] = self.final_topk
552
- config['multi_retriever_setting']['device'] = config['device']
553
- self.reranker = get_reranker(config['multi_retriever_setting'])
559
+ if self.merge_method == "rerank":
560
+ config["multi_retriever_setting"]["rerank_topk"] = self.final_topk
561
+ config["multi_retriever_setting"]["device"] = config["device"]
562
+ self.reranker = get_reranker(config["multi_retriever_setting"])
554
563
 
555
564
  def load_all_retriever(self, config):
556
565
  retriever_config_list = config["multi_retriever_setting"]["retriever_list"]
@@ -586,7 +595,7 @@ class MultiRetrieverRouter:
586
595
  try:
587
596
  model_config = AutoConfig.from_pretrained(retrieval_model_path)
588
597
  arch = model_config.architectures[0]
589
- print("arch: ",arch)
598
+ print("arch: ", arch)
590
599
  if "clip" in arch.lower():
591
600
  retriever = MultiModalRetriever(retriever_config, corpus)
592
601
  else:
@@ -609,11 +618,11 @@ class MultiRetrieverRouter:
609
618
  for _item in item:
610
619
  _item["source"] = retrieval_method
611
620
  _item["corpus_path"] = corpus_path
612
- _item['is_multimodal'] = is_multimodal
621
+ _item["is_multimodal"] = is_multimodal
613
622
  else:
614
623
  item["source"] = retrieval_method
615
624
  item["corpus_path"] = corpus_path
616
- item['is_multimodal'] = is_multimodal
625
+ item["is_multimodal"] = is_multimodal
617
626
  return result
618
627
 
619
628
  def _search_or_batch_search(self, query: Union[str, list], target_modal, num, return_score, method, retriever_list):
@@ -626,7 +635,7 @@ class MultiRetrieverRouter:
626
635
  def process_retriever(retriever):
627
636
  is_multimodal = isinstance(retriever, MultiModalRetriever)
628
637
  params = {"query": query, "return_score": return_score}
629
-
638
+
630
639
  if is_multimodal:
631
640
  params["target_modal"] = target_modal
632
641
 
@@ -645,7 +654,9 @@ class MultiRetrieverRouter:
645
654
  return result, score
646
655
 
647
656
  with ThreadPoolExecutor(max_workers=4) as executor:
648
- future_to_retriever = {executor.submit(process_retriever, retriever): retriever for retriever in retriever_list}
657
+ future_to_retriever = {
658
+ executor.submit(process_retriever, retriever): retriever for retriever in retriever_list
659
+ }
649
660
  for future in as_completed(future_to_retriever):
650
661
  try:
651
662
  result, score = future.result()
@@ -661,7 +672,6 @@ class MultiRetrieverRouter:
661
672
  else:
662
673
  return result_list
663
674
 
664
-
665
675
  def reorder(self, result_list, score_list, retriever_list):
666
676
  """
667
677
  batch_search:
@@ -670,7 +680,7 @@ class MultiRetrieverRouter:
670
680
 
671
681
  navie search:
672
682
  original result like: [bm25-d1, bm25-d2, e5-d1, e5-d2]
673
-
683
+
674
684
  """
675
685
 
676
686
  retriever_num = len(retriever_list)
@@ -705,8 +715,8 @@ class MultiRetrieverRouter:
705
715
  for query_idx, query_doc_list in enumerate(result_list):
706
716
  exist_id = set()
707
717
  for doc_idx, doc in enumerate(query_doc_list):
708
- if doc['id'] not in exist_id:
709
- exist_id.add(doc['id'])
718
+ if doc["id"] not in exist_id:
719
+ exist_id.add(doc["id"])
710
720
  else:
711
721
  query_doc_list.remove(doc)
712
722
  if score_list != []:
@@ -726,14 +736,14 @@ class MultiRetrieverRouter:
726
736
  else:
727
737
  result_list, score_list = self.rrf_merge(result_list, num, k=60)
728
738
  return result_list, score_list
729
- elif self.merge_method == 'rerank':
739
+ elif self.merge_method == "rerank":
730
740
  if isinstance(result_list[0], dict):
731
741
  query, result_list, score_list = [query], [result_list], [score_list]
732
742
  # parse the result of multimodal corpus
733
743
  for item_result in result_list:
734
744
  for item in item_result:
735
- if item['is_multimodal']:
736
- item['contents'] = item['text']
745
+ if item["is_multimodal"]:
746
+ item["contents"] = item["text"]
737
747
  # rerank all docs
738
748
  print(result_list)
739
749
  result_list, score_list = self.reranker.rerank(query, result_list, topk=num)
@@ -793,50 +803,75 @@ class MultiRetrieverRouter:
793
803
  # query: str or PIL.Image
794
804
  # judge query type: text or image
795
805
  if judge_image(query):
796
- retriever_list = [retriever for retriever in self.retriever_list if isinstance(retriever, MultiModalRetriever)]
806
+ retriever_list = [
807
+ retriever for retriever in self.retriever_list if isinstance(retriever, MultiModalRetriever)
808
+ ]
797
809
  else:
798
810
  retriever_list = self.retriever_list
799
- if target_modal == 'image':
811
+ if target_modal == "image":
800
812
  # remove text retriever
801
813
  retriever_list = [retriever for retriever in retriever_list if isinstance(retriever, MultiModalRetriever)]
802
814
 
803
- return self._search_or_batch_search(query, target_modal, num, return_score, method="search", retriever_list=retriever_list)
815
+ return self._search_or_batch_search(
816
+ query, target_modal, num, return_score, method="search", retriever_list=retriever_list
817
+ )
804
818
 
805
819
  def batch_search(self, query, target_modal="text", num: Union[list, int, None] = None, return_score=False):
806
820
  # judge query type: text or image
807
821
  if not isinstance(query, list):
808
822
  query = [query]
809
- if target_modal == 'image':
810
- self._retriever_list = [retriever for retriever in self.retriever_list if isinstance(retriever, MultiModalRetriever)]
823
+ if target_modal == "image":
824
+ self._retriever_list = [
825
+ retriever for retriever in self.retriever_list if isinstance(retriever, MultiModalRetriever)
826
+ ]
811
827
  else:
812
828
  self._retriever_list = self.retriever_list
813
829
  query_type_list = [judge_image(q) for q in query]
814
830
  if all(query_type_list):
815
831
  # all query is image
816
- if self.merge_method == 'rerank':
817
- warnings.warn('merge_method is rerank, but all query is image, use default method `concat` instead')
818
- self.merge_method = 'concat'
819
- retriever_list = [retriever for retriever in self._retriever_list if isinstance(retriever, MultiModalRetriever)]
820
-
821
- return self._search_or_batch_search(query, target_modal, num, return_score, method="batch_search", retriever_list=retriever_list)
832
+ if self.merge_method == "rerank":
833
+ warnings.warn("merge_method is rerank, but all query is image, use default method `concat` instead")
834
+ self.merge_method = "concat"
835
+ retriever_list = [
836
+ retriever for retriever in self._retriever_list if isinstance(retriever, MultiModalRetriever)
837
+ ]
838
+
839
+ return self._search_or_batch_search(
840
+ query, target_modal, num, return_score, method="batch_search", retriever_list=retriever_list
841
+ )
822
842
  elif all([not t for t in query_type_list]):
823
843
  # all query is text
824
844
  # if exist text retriever, don't use mm retriever for text-text search
825
845
  if any([isinstance(retriever, BaseTextRetriever) for retriever in self._retriever_list]):
826
- self._retriever_list = [retriever for retriever in self._retriever_list if not isinstance(retriever, MultiModalRetriever)]
827
- return self._search_or_batch_search(query, target_modal, num, return_score, method="batch_search", retriever_list=self._retriever_list)
846
+ self._retriever_list = [
847
+ retriever for retriever in self._retriever_list if not isinstance(retriever, MultiModalRetriever)
848
+ ]
849
+ return self._search_or_batch_search(
850
+ query, target_modal, num, return_score, method="batch_search", retriever_list=self._retriever_list
851
+ )
828
852
  else:
829
853
  # query list is the mix of image and text
830
- if self.merge_method == 'rerank':
831
- warnings.warn('merge_method is rerank, but some query is image, use default method `concat` instead')
832
- self.merge_method = 'concat'
854
+ if self.merge_method == "rerank":
855
+ warnings.warn("merge_method is rerank, but some query is image, use default method `concat` instead")
856
+ self.merge_method = "concat"
833
857
  image_query_idx = [i for i, t in enumerate(query_type_list) if t]
834
858
  image_query_list = [query[i] for i in image_query_idx]
835
859
  text_query_list = [q for q in query if q not in image_query_list]
836
860
 
837
- text_output = self._search_or_batch_search(text_query_list, target_modal, num, return_score, method="batch_search", retriever_list=self._retriever_list)
838
- retriever_list = [retriever for retriever in self._retriever_list if isinstance(retriever, MultiModalRetriever)]
839
- image_output = self._search_or_batch_search(text_query_list, target_modal, num, return_score, method="batch_search", retriever_list=retriever_list)
861
+ text_output = self._search_or_batch_search(
862
+ text_query_list,
863
+ target_modal,
864
+ num,
865
+ return_score,
866
+ method="batch_search",
867
+ retriever_list=self._retriever_list,
868
+ )
869
+ retriever_list = [
870
+ retriever for retriever in self._retriever_list if isinstance(retriever, MultiModalRetriever)
871
+ ]
872
+ image_output = self._search_or_batch_search(
873
+ text_query_list, target_modal, num, return_score, method="batch_search", retriever_list=retriever_list
874
+ )
840
875
 
841
876
  # merge text output and image output
842
877
  if return_score:
@@ -854,7 +889,7 @@ class MultiRetrieverRouter:
854
889
  else:
855
890
  final_result.append(image_result[image_idx])
856
891
  final_score.append(image_score[image_idx])
857
- image_idx += 1
892
+ image_idx += 1
858
893
  return final_result, final_score
859
894
  else:
860
895
  final_result = []
@@ -866,6 +901,5 @@ class MultiRetrieverRouter:
866
901
  text_idx += 1
867
902
  else:
868
903
  final_result.append(image_result[image_idx])
869
- image_idx += 1
904
+ image_idx += 1
870
905
  return final_result
871
-
@@ -8,6 +8,8 @@ import re
8
8
  import langid
9
9
  from transformers import AutoTokenizer, AutoModel, AutoConfig
10
10
 
11
+ _has_printed_instruction = False # trigger instruction print once
12
+
11
13
  def convert_numpy(obj: Union[Dict, list, np.ndarray, np.generic]) -> Any:
12
14
  """Recursively convert numpy objects in nested dictionaries or lists to native Python types."""
13
15
  if isinstance(obj, dict):
@@ -101,6 +103,7 @@ def parse_query(model_name, query_list, instruction=None, is_query=True):
101
103
  """
102
104
  processing query for different encoders
103
105
  """
106
+ global _has_printed_instruction
104
107
 
105
108
  if isinstance(query_list, str):
106
109
  query_list = [query_list]
@@ -109,10 +112,14 @@ def parse_query(model_name, query_list, instruction=None, is_query=True):
109
112
  instruction = instruction.strip() + " "
110
113
  else:
111
114
  instruction = set_default_instruction(model_name, is_query=is_query, is_zh=judge_zh(query_list[0]))
112
- print(f"Use `{instruction}` as retreival instruction")
113
- if instruction == "":
114
- warnings.warn('Instruction is not set')
115
-
115
+
116
+ if not _has_printed_instruction:
117
+ if instruction == "":
118
+ warnings.warn('Instruction is not set')
119
+ else:
120
+ print(f"Use `{instruction}` as retreival instruction")
121
+ _has_printed_instruction = True
122
+
116
123
  query_list = [instruction + query for query in query_list]
117
124
 
118
125
  return query_list
@@ -1,4 +1,5 @@
1
1
  import os
2
+ import re
2
3
  import json
3
4
  import importlib
4
5
  from transformers import AutoConfig
@@ -173,3 +174,10 @@ def hash_object(o) -> str:
173
174
  dill.dump(o, buffer)
174
175
  m.update(buffer.getbuffer())
175
176
  return base58.b58encode(m.digest()).decode()
177
+
178
+ def extract_between(text: str, start_tag: str, end_tag: str):
179
+ pattern = re.escape(start_tag) + r"(.*?)" + re.escape(end_tag)
180
+ matches = re.findall(pattern, text, flags=re.DOTALL)
181
+ if matches:
182
+ return matches[-1].strip()
183
+ return None
@@ -0,0 +1 @@
1
+ __version__ = "0.1.4dev20250321"
@@ -1,3 +1,20 @@
1
+ Metadata-Version: 2.1
2
+ Name: flashrag-dev
3
+ Version: 0.1.4.dev20250321
4
+ Summary: A library for efficient Retrieval-Augmented Generation research
5
+ Home-page: https://github.com/RUC-NLPIR/FlashRAG
6
+ Author: Jiajie Jin, Yutao Zhu, Chenghao Zhang, Xinyu Yang, Zhicheng Dou
7
+ Author-email: jinjiajie@ruc.edu.cn
8
+ License: MIT License
9
+ Requires-Python: >=3.9
10
+ Description-Content-Type: text/markdown
11
+ Provides-Extra: core
12
+ Provides-Extra: retriever
13
+ Provides-Extra: generator
14
+ Provides-Extra: multimodal
15
+ Provides-Extra: full
16
+ License-File: LICENSE
17
+
1
18
  # <div align="center">⚡FlashRAG: A Python Toolkit for Efficient RAG Research<div>
2
19
  \[ English | [中文](README_zh.md) \]
3
20
  <div align="center">
@@ -24,7 +41,7 @@
24
41
  </h4>
25
42
 
26
43
 
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.
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.
28
45
 
29
46
  <p align="center">
30
47
  <img src="asset/framework.jpg">
@@ -58,7 +75,7 @@ https://github.com/user-attachments/assets/8ca00873-5df2-48a7-b853-89e7b18bc6e9
58
75
 
59
76
  - **Comprehensive Benchmark Datasets**: A collection of 36 pre-processed RAG benchmark datasets to test and validate RAG models' performances.
60
77
 
61
- - **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.
62
79
 
63
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.
64
81
 
@@ -79,6 +96,8 @@ FlashRAG is still under development and there are many issues and room for impro
79
96
  - [ ] Enhance code adaptability and readability
80
97
 
81
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
+
82
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!
83
102
 
84
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!
@@ -465,7 +484,7 @@ In each category, we have implemented corresponding common pipelines. Some pipel
465
484
  <td>Ranking and merging generated results based on each document</td>
466
485
  </tr>
467
486
  <tr>
468
- <td rowspan="5">Loop</td>
487
+ <td rowspan="6">Loop</td>
469
488
  <td>Iterative Pipeline</td>
470
489
  <td>Alternating retrieval and generation</td>
471
490
  </tr>
@@ -485,6 +504,10 @@ In each category, we have implemented corresponding common pipelines. Some pipel
485
504
  <td>IRCoT Pipeline</td>
486
505
  <td>Integrate retrieval process with CoT</td>
487
506
  </tr>
507
+ <tr>
508
+ <td>Reasoning Pipeline</td>
509
+ <td>Reasoning with retrieval</td>
510
+ </tr>
488
511
  </tbody>
489
512
  </table>
490
513
 
@@ -565,6 +588,8 @@ It’s important to note that, to ensure consistency, we have utilized a uniform
565
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 | |
566
589
  | [IRCoT](https://aclanthology.org/2023.acl-long.557.pdf) | Loop | 33.3 | 56.9 | 41.5 | 32.4 | 45.6 | 20.7 | |
567
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
+
568
593
 
569
594
 
570
595
  ## :notebook: Supporting Datasets & Document Corpus
@@ -29,6 +29,7 @@ flashrag/pipeline/active_pipeline.py
29
29
  flashrag/pipeline/branching_pipeline.py
30
30
  flashrag/pipeline/mm_pipeline.py
31
31
  flashrag/pipeline/pipeline.py
32
+ flashrag/pipeline/reasoning_pipeline.py
32
33
  flashrag/pipeline/replug_utils.py
33
34
  flashrag/prompt/__init__.py
34
35
  flashrag/prompt/base_prompt.py
@@ -1 +0,0 @@
1
- __version__ = "0.1.4dev20250319"