flashrag-dev 0.1.1__tar.gz → 0.1.2.dev20241101__tar.gz

{flashrag_dev-0.1.1/flashrag_dev.egg-info → flashrag-dev-0.1.2.dev20241101}/PKG-INFO

@@ -1,75 +1,18 @@
 Metadata-Version: 2.1
 Name: flashrag-dev
-Version: 0.1.1
+Version: 0.1.2.dev20241101
 Summary: A library for efficient Retrieval-Augmented Generation research
 Home-page: https://github.com/RUC-NLPIR/FlashRAG
 Author: Jiajie Jin, Yutao Zhu, Chenghao Zhang, Xinyu Yang, Zhicheng Dou
 Author-email: jinjiajie@ruc.edu.cn
 License: MIT License
-Requires-Python: >=3.8
+Requires-Python: >=3.9
 Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: datasets
-Requires-Dist: base58
-Requires-Dist: nltk
-Requires-Dist: numpy
-Requires-Dist: openai
-Requires-Dist: peft
-Requires-Dist: PyYAML
-Requires-Dist: rank_bm25
-Requires-Dist: rouge
-Requires-Dist: spacy>=3.6.1
-Requires-Dist: tiktoken
-Requires-Dist: torch
-Requires-Dist: tqdm
-Requires-Dist: transformers>=4.40.0
-Requires-Dist: bm25s[full]
-Requires-Dist: fschat
-Requires-Dist: streamlit
 Provides-Extra: core
-Requires-Dist: datasets; extra == "core"
-Requires-Dist: base58; extra == "core"
-Requires-Dist: nltk; extra == "core"
-Requires-Dist: numpy; extra == "core"
-Requires-Dist: openai; extra == "core"
-Requires-Dist: peft; extra == "core"
-Requires-Dist: PyYAML; extra == "core"
-Requires-Dist: rank_bm25; extra == "core"
-Requires-Dist: rouge; extra == "core"
-Requires-Dist: spacy>=3.6.1; extra == "core"
-Requires-Dist: tiktoken; extra == "core"
-Requires-Dist: torch; extra == "core"
-Requires-Dist: tqdm; extra == "core"
-Requires-Dist: transformers>=4.40.0; extra == "core"
-Requires-Dist: bm25s[full]; extra == "core"
-Requires-Dist: fschat; extra == "core"
-Requires-Dist: streamlit; extra == "core"
 Provides-Extra: retriever
-Requires-Dist: pyserini; extra == "retriever"
-Requires-Dist: sentence-transformers>=3.0.1; extra == "retriever"
 Provides-Extra: generator
-Requires-Dist: vllm>=0.4.1; extra == "generator"
 Provides-Extra: full
-Requires-Dist: datasets; extra == "full"
-Requires-Dist: base58; extra == "full"
-Requires-Dist: nltk; extra == "full"
-Requires-Dist: numpy; extra == "full"
-Requires-Dist: openai; extra == "full"
-Requires-Dist: peft; extra == "full"
-Requires-Dist: PyYAML; extra == "full"
-Requires-Dist: rank_bm25; extra == "full"
-Requires-Dist: rouge; extra == "full"
-Requires-Dist: spacy>=3.6.1; extra == "full"
-Requires-Dist: tiktoken; extra == "full"
-Requires-Dist: torch; extra == "full"
-Requires-Dist: tqdm; extra == "full"
-Requires-Dist: transformers>=4.40.0; extra == "full"
-Requires-Dist: bm25s[full]; extra == "full"
-Requires-Dist: fschat; extra == "full"
-Requires-Dist: streamlit; extra == "full"
-Requires-Dist: pyserini; extra == "full"
-Requires-Dist: sentence-transformers>=3.0.1; extra == "full"
-Requires-Dist: vllm>=0.4.1; extra == "full"
+License-File: LICENSE
 
 # <div align="center">⚡FlashRAG: A Python Toolkit for Efficient RAG Research<div>
 
@@ -94,7 +37,7 @@ Requires-Dist: vllm>=0.4.1; extra == "full"
 </p>
 
 </h4>
-FlashRAG is a Python toolkit for the reproduction and development of Retrieval Augmented Generation (RAG) research. Our toolkit includes 32 pre-processed benchmark RAG datasets and 15 state-of-the-art RAG algorithms. 
+FlashRAG is a Python toolkit for the reproduction and development of Retrieval Augmented Generation (RAG) research. Our toolkit includes 33 pre-processed benchmark RAG datasets and 15 state-of-the-art RAG algorithms. 
 
 <p align="center">
 <img src="asset/framework.jpg">
@@ -110,7 +53,7 @@ With FlashRAG and provided resources, you can effortlessly reproduce existing SO
 
 - **Extensive and Customizable Framework**: Includes essential components for RAG scenarios such as retrievers, rerankers, generators, and compressors, allowing for flexible assembly of complex pipelines.
 
-- **Comprehensive Benchmark Datasets**: A collection of 32 pre-processed RAG benchmark datasets to test and validate RAG models' performances.
+- **Comprehensive Benchmark Datasets**: A collection of 33 pre-processed RAG benchmark datasets to test and validate RAG models' performances.
 
 - **Pre-implemented Advanced RAG Algorithms**: Features 15 advancing RAG algorithms with reported results, based on our framework. Easily reproducing results under different settings.
 
@@ -123,7 +66,6 @@ With FlashRAG and provided resources, you can effortlessly reproduce existing SO
 FlashRAG is still under development and there are many issues and room for improvement. We will continue to update. And we also sincerely welcome contributions on this open-source toolkit.
 
 - [x] Support OpenAI models
-- [ ] Support Claude and Gemini models
 - [x] Provdide instructions for each component
 - [x] Integrate sentence Transformers
 - [ ] Inlcude more RAG approaches
@@ -132,6 +74,11 @@ FlashRAG is still under development and there are many issues and room for impro
 
 
 ## :page_with_curl: Changelog
+[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.
+
+[24/10/13] A new in-domain dataset and corpus - [DomainRAG](https://arxiv.org/pdf/2406.05654) have been added to the dataset. The dataset is based on the internal enrollment data of Renmin University of China, covering seven types of tasks, which can be used for conducting domain-specific RAG testing.
+
+[24/09/24] We have released a version based on the MindSpore framework that supports Chinese hardware platforms. Please refer to [FlashRAG MindSpore](https://github.com/RUC-NLPIR/FlashRAG-MindSpore) for details.
 
 [24/09/18] Due to the complexity and limitations of installing Pyserini in certain environments, we have introduced a lightweight `BM25s` package as an alternative (faster and easier to use). The retriever based on Pyserini will be deprecated in future versions. To use retriever with `bm25s`, just set `bm25_backend` to `bm25s` in config.
 
@@ -165,7 +112,7 @@ FlashRAG is still under development and there are many issues and room for impro
 
 To get started with FlashRAG, you can simply install it with pip:
 ```base
-pip install flashrag[core]
+pip install flashrag-dev[core]
 ```
 
 Or you can clone it from Github and install (requires Python 3.9+): 
@@ -554,6 +501,7 @@ Below is the list of datasets along with the corresponding sample sizes:
 | multi-hop QA              | 2WikiMultiHopQA | wiki             | 15,000    | 12,576  | /      |
 | multi-hop QA              | Musique         | wiki             | 19,938    | 2,417   | /      |
 | multi-hop QA              | Bamboogle       | wiki             | /         | /       | 125    |
+| multi-hop QA              | StrategyQA      | wiki             | 2290      | /       | /
 | Long-form QA              | ASQA            | wiki             | 4,353     | 948     | /      |
 | Long-form QA              | ELI5            | Reddit           | 272,634   | 1,507   | /      |
 | Open-Domain Summarization | WikiASP         | wiki             | 300,636   | 37,046  | 37,368 |
@@ -568,6 +516,7 @@ Below is the list of datasets along with the corresponding sample sizes:
 | Entity Linking            | WNED            | Wiki             | /         | 8,995   | /      |
 | Slot Filling              | T-REx           | DBPedia          | 2,284,168 | 5,000   | /      |
 | Slot Filling              | Zero-shot RE    | wiki             | 147,909   | 3,724   | /      |
+| In-domain QA| DomainRAG | Web pages of RUC| / | / | 485|
 
 ### Document Corpus
 

{flashrag_dev-0.1.1 → flashrag-dev-0.1.2.dev20241101}/README.md

@@ -21,7 +21,7 @@
 </p>
 
 </h4>
-FlashRAG is a Python toolkit for the reproduction and development of Retrieval Augmented Generation (RAG) research. Our toolkit includes 32 pre-processed benchmark RAG datasets and 15 state-of-the-art RAG algorithms. 
+FlashRAG is a Python toolkit for the reproduction and development of Retrieval Augmented Generation (RAG) research. Our toolkit includes 33 pre-processed benchmark RAG datasets and 15 state-of-the-art RAG algorithms. 
 
 <p align="center">
 <img src="asset/framework.jpg">
@@ -37,7 +37,7 @@ With FlashRAG and provided resources, you can effortlessly reproduce existing SO
 
 - **Extensive and Customizable Framework**: Includes essential components for RAG scenarios such as retrievers, rerankers, generators, and compressors, allowing for flexible assembly of complex pipelines.
 
-- **Comprehensive Benchmark Datasets**: A collection of 32 pre-processed RAG benchmark datasets to test and validate RAG models' performances.
+- **Comprehensive Benchmark Datasets**: A collection of 33 pre-processed RAG benchmark datasets to test and validate RAG models' performances.
 
 - **Pre-implemented Advanced RAG Algorithms**: Features 15 advancing RAG algorithms with reported results, based on our framework. Easily reproducing results under different settings.
 
@@ -50,7 +50,6 @@ With FlashRAG and provided resources, you can effortlessly reproduce existing SO
 FlashRAG is still under development and there are many issues and room for improvement. We will continue to update. And we also sincerely welcome contributions on this open-source toolkit.
 
 - [x] Support OpenAI models
-- [ ] Support Claude and Gemini models
 - [x] Provdide instructions for each component
 - [x] Integrate sentence Transformers
 - [ ] Inlcude more RAG approaches
@@ -59,6 +58,11 @@ FlashRAG is still under development and there are many issues and room for impro
 
 
 ## :page_with_curl: Changelog
+[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.
+
+[24/10/13] A new in-domain dataset and corpus - [DomainRAG](https://arxiv.org/pdf/2406.05654) have been added to the dataset. The dataset is based on the internal enrollment data of Renmin University of China, covering seven types of tasks, which can be used for conducting domain-specific RAG testing.
+
+[24/09/24] We have released a version based on the MindSpore framework that supports Chinese hardware platforms. Please refer to [FlashRAG MindSpore](https://github.com/RUC-NLPIR/FlashRAG-MindSpore) for details.
 
 [24/09/18] Due to the complexity and limitations of installing Pyserini in certain environments, we have introduced a lightweight `BM25s` package as an alternative (faster and easier to use). The retriever based on Pyserini will be deprecated in future versions. To use retriever with `bm25s`, just set `bm25_backend` to `bm25s` in config.
 
@@ -92,7 +96,7 @@ FlashRAG is still under development and there are many issues and room for impro
 
 To get started with FlashRAG, you can simply install it with pip:
 ```base
-pip install flashrag[core]
+pip install flashrag-dev[core]
 ```
 
 Or you can clone it from Github and install (requires Python 3.9+): 
@@ -481,6 +485,7 @@ Below is the list of datasets along with the corresponding sample sizes:
 | multi-hop QA              | 2WikiMultiHopQA | wiki             | 15,000    | 12,576  | /      |
 | multi-hop QA              | Musique         | wiki             | 19,938    | 2,417   | /      |
 | multi-hop QA              | Bamboogle       | wiki             | /         | /       | 125    |
+| multi-hop QA              | StrategyQA      | wiki             | 2290      | /       | /
 | Long-form QA              | ASQA            | wiki             | 4,353     | 948     | /      |
 | Long-form QA              | ELI5            | Reddit           | 272,634   | 1,507   | /      |
 | Open-Domain Summarization | WikiASP         | wiki             | 300,636   | 37,046  | 37,368 |
@@ -495,6 +500,7 @@ Below is the list of datasets along with the corresponding sample sizes:
 | Entity Linking            | WNED            | Wiki             | /         | 8,995   | /      |
 | Slot Filling              | T-REx           | DBPedia          | 2,284,168 | 5,000   | /      |
 | Slot Filling              | Zero-shot RE    | wiki             | 147,909   | 3,724   | /      |
+| In-domain QA| DomainRAG | Web pages of RUC| / | / | 485|
 
 ### Document Corpus
 

flashrag-dev-0.1.2.dev20241101/flashrag/config/basic_config.yaml

@@ -0,0 +1,104 @@
+# ------------------------------------------------Global Paths------------------------------------------------#
+# Paths to various models
+model2path:
+  e5: "intfloat/e5-base-v2"
+  bge: "intfloat/e5-base-v2"
+  contriever: "facebook/contriever"
+  llama2-7B-chat: "meta-llama/Llama-2-7b-chat-hf"
+  llama2-7B: "meta-llama/Llama-2-7b-hf"
+  llama2-13B: "meta-llama/Llama-2-13b-hf"
+  llama2-13B-chat: "meta-llama/Llama-2-13b-chat-hf"
+  
+# Pooling methods for each embedding model
+model2pooling:
+  e5: "mean"
+  bge: "cls"
+  contriever: "mean"
+  jina: 'mean'
+  dpr: cls
+
+# Indexes path for retrieval models
+method2index:
+  e5: ~
+  bm25: ~
+  contriever: ~
+
+# ------------------------------------------------Environment Settings------------------------------------------------#
+# Directory paths for data and outputs
+data_dir: "dataset/"
+save_dir: "output/"
+
+gpu_id: "0,1,2,3"
+dataset_name: "nq" # name of the dataset in data_dir
+split: ["test"]  # dataset split to load (e.g. train,dev,test)
+
+# Sampling configurations for testing
+test_sample_num: ~  # number of samples to test (only work in dev/test split), if None, test all samples
+random_sample: False # whether to randomly sample the test samples
+
+# Seed for reproducibility
+seed: 2024
+
+# Whether save intermediate data
+save_intermediate_data: True
+save_note: 'experiment'
+
+# -------------------------------------------------Retrieval Settings------------------------------------------------#
+# If set the name, the model path will be find in global paths
+retrieval_method: "e5"  # name or path of the retrieval model. 
+retrieval_model_path: ~ # path to the retrieval model
+index_path: ~ # set automatically if not provided. 
+faiss_gpu: False # whether use gpu to hold index
+corpus_path: ~  # path to corpus in '.jsonl' format that store the documents
+
+instruction: ~ # instruction for the retrieval model
+retrieval_topk: 5 # number of retrieved documents
+retrieval_batch_size: 256  # batch size for retrieval
+retrieval_use_fp16: True  # whether to use fp16 for retrieval model
+retrieval_query_max_length: 128  # max length of the query
+save_retrieval_cache: False # whether to save the retrieval cache
+use_retrieval_cache: False # whether to use the retrieval cache
+retrieval_cache_path: ~ # path to the retrieval cache
+retrieval_pooling_method: ~ # set automatically if not provided
+bm25_backend: bm25s # pyserini, bm25s
+use_sentence_transformer: False
+
+use_reranker: False # whether to use reranker
+rerank_model_name: ~ # same as retrieval_method
+rerank_model_path: ~ # path to reranker model, path will be automatically find in `model2path`
+rerank_pooling_method: ~
+rerank_topk: 5  # number of remain documents after reranking
+rerank_max_length: 512 
+rerank_batch_size: 256 # batch size for reranker
+rerank_use_fp16: True
+
+# -------------------------------------------------Generator Settings------------------------------------------------#
+framework: fschat # inference frame work of LLM, supporting: 'hf','vllm','fschat', 'openai'
+generator_model: "llama3-8B-instruct" # name or path of the generator model
+# setting for openai model, only valid in openai framework
+openai_setting:
+  api_key: ~
+  base_url: ~
+
+generator_model_path: ~
+generator_max_input_len: 1024  # max length of the input
+generator_batch_size: 4 # batch size for generation, invalid for vllm
+generation_params:  
+  #do_sample: false
+  max_tokens: 32
+  #temperature: 1.0
+  #top_p: 1.0
+use_fid: False # whether to use FID, only valid in encoder-decoder model
+gpu_memory_utilization: 0.85 # ratio of gpu's memory usage for generator
+
+# -------------------------------------------------Evaluation Settings------------------------------------------------#
+# Metrics to evaluate the result
+metrics: ['em','f1','acc','precision','recall','input_tokens'] 
+# Specify setting for metric, will be called within certain metrics
+metric_setting: 
+  retrieval_recall_topk: 5
+  tokenizer_name: 'gpt-4'
+save_metric_score: True #　whether to save the metric score into txt file
+
+
+

{flashrag_dev-0.1.1 → flashrag-dev-0.1.2.dev20241101}/flashrag/config/config.py

@@ -158,6 +158,8 @@ class Config:
             refiner_model = self.final_config["refiner_name"]
             if "refiner_model_path" not in self.final_config or self.final_config["refiner_model_path"] is None:
                 self.final_config["refiner_model_path"] = model2path.get(refiner_model, None)
+        if 'instruction' not in self.final_config or self.final_config['instruction'] is None:
+            self.final_config['instruction'] = ""
 
         # set model path in metric setting
         metric_setting = self.final_config["metric_setting"]

{flashrag_dev-0.1.1 → flashrag-dev-0.1.2.dev20241101}/flashrag/generator/generator.py

@@ -24,7 +24,7 @@ class BaseGenerator:
         self.batch_size = config["generator_batch_size"]
         self.device = config["device"]
         self.gpu_num = torch.cuda.device_count()
-
+        self.config = config
         self.generation_params = config["generation_params"]
 
     def generate(self, input_list: list) -> List[str]:
@@ -217,8 +217,10 @@ class VLLMGenerator(BaseGenerator):
         generation_params = deepcopy(self.generation_params)
         generation_params.update(params)
         if "do_sample" in generation_params:
-            generation_params.pop("do_sample")
-
+            do_sample_flag = generation_params.pop("do_sample")
+            if not do_sample_flag:
+                generation_params['temperature'] = 0
+        generation_params['seed'] = self.config['seed']
         max_tokens = params.pop("max_tokens", None) or params.pop(
             "max_new_tokens", None
         )

{flashrag_dev-0.1.1 → flashrag-dev-0.1.2.dev20241101}/flashrag/pipeline/active_pipeline.py

@@ -18,7 +18,6 @@ class IterativePipeline(BasicPipeline):
 
     def run(self, dataset, do_eval=True, pred_process_fun=None):
         questions = dataset.question
-
         # run in batch
         past_generation_result = []  # list of N items
         for iter_idx in range(self.iter_num):

{flashrag_dev-0.1.1 → flashrag-dev-0.1.2.dev20241101}/flashrag/pipeline/branching_pipeline.py

@@ -34,7 +34,6 @@ class REPLUGPipeline(BasicPipeline):
         from flashrag.pipeline.replug_utils import REPLUGLogitsProcessor
 
         input_query = dataset.question
-
         retrieval_results, doc_scores = self.retriever.batch_search(input_query, return_score=True)
         dataset.update_output("retrieval_result", retrieval_results)
         dataset.update_output("doc_scores", doc_scores)

{flashrag_dev-0.1.1 → flashrag-dev-0.1.2.dev20241101}/flashrag/pipeline/pipeline.py

@@ -84,7 +84,6 @@ class SequentialPipeline(BasicPipeline):
 
     def run(self, dataset, do_eval=True, pred_process_fun=None):
         input_query = dataset.question
-
         retrieval_results = self.retriever.batch_search(input_query)
         dataset.update_output("retrieval_result", retrieval_results)
 

{flashrag_dev-0.1.1 → flashrag-dev-0.1.2.dev20241101}/flashrag/prompt/base_prompt.py

@@ -1,5 +1,5 @@
 from transformers import AutoTokenizer, AutoConfig
-
+import tiktoken
 
 class PromptTemplate:
     placeholders = ["reference", "question"]
@@ -14,6 +14,7 @@ class PromptTemplate:
 
         self.config = config
         self.is_openai = config["framework"] == "openai"
+        self.max_input_len = config['generator_max_input_len']
         if not self.is_openai:
             self.generator_path = config["generator_model_path"]
             model_config = AutoConfig.from_pretrained(self.generator_path, trust_remote_code=True)
@@ -21,10 +22,11 @@ class PromptTemplate:
             self.is_chat = False
             if "chat" in model_name or "instruct" in model_name:
                 self.is_chat = True
-                self.tokenizer = AutoTokenizer.from_pretrained(self.generator_path, trust_remote_code=True)
+            self.tokenizer = AutoTokenizer.from_pretrained(self.generator_path, trust_remote_code=True)
         else:
             self.is_chat = True
             self.enable_chat = True
+            self.tokenizer = tiktoken.encoding_for_model("gpt-4o")
 
         if len(system_prompt) == 0 and len(user_prompt) == 0:
             system_prompt = self.base_system_prompt
@@ -48,26 +50,61 @@ class PromptTemplate:
             if not flag and holder != "reference":
                 assert False
 
+    def truncate_prompt(self, prompt):
+        if self.is_openai:
+            truncated_messages = []
+            total_tokens = 0
+            assert isinstance(prompt, list)
+            for message in prompt:
+                role_content = message['content']
+                encoded_message = self.tokenizer.encode(role_content)
+
+                if total_tokens + len(encoded_message) <= self.max_input_len:
+                    truncated_messages.append(message)
+                    total_tokens += len(encoded_message)
+                else:
+                    print(f"The input text length is greater than the maximum length ({total_tokens + len(encoded_message)} > {self.max_input_len}) and has been truncated!")
+                    remaining_tokens = self.max_input_len - total_tokens
+                    truncated_message = self.encoding.decode(encoded_message[:remaining_tokens])
+                    message['content'] = truncated_message
+                    truncated_messages.append(message)
+                    break
+
+            return truncated_messages
+
+        else:
+            assert isinstance(prompt, str)
+            tokenized_prompt = self.tokenizer(prompt, truncation=False, return_tensors="pt").input_ids[0]
+
+            if len(tokenized_prompt) > self.max_input_len:
+                print(f"The input text length is greater than the maximum length ({len(tokenized_prompt)} > {self.max_input_len}) and has been truncated!")
+                half = int(self.max_input_len / 2)
+                prompt = self.tokenizer.decode(tokenized_prompt[:half], skip_special_tokens=True) + \
+                        self.tokenizer.decode(tokenized_prompt[-half:], skip_special_tokens=True)
+            return prompt
+
+
+
     def get_string(self, question=None, retrieval_result=None, formatted_reference=None, previous_gen=None, messages=None, **params):
         if messages is not None:
             if isinstance(messages, str):
-                return messages
+                return self.truncate_prompt(messages)
             if self.is_chat and self.enable_chat:
                 if self.is_openai:
                     for item in input:
                         if item["role"] == "system":
                             item["role"] = "assistant"
-                    return messages
+                    self.truncate_prompt(messages)
                 else:
                     prompt = self.tokenizer.apply_chat_template(
                         messages, tokenize=False, add_generation_prompt=True
                     )
-                    return prompt
+                    return self.truncate_prompt(prompt)
             else:
                 prompt = "\n\n".join(
                     [message['content'] for message in messages if message['content']]
                 )
-                return prompt
+                return self.truncate_prompt(prompt)
 
         if formatted_reference is None:
             if retrieval_result is not None:
@@ -99,7 +136,7 @@ class PromptTemplate:
         if previous_gen is not None and previous_gen not in ["", " "] and self.is_openai is False:
             input += previous_gen
 
-        return input
+        return self.truncate_prompt(input)
 
     def get_string_with_varying_examplars(
         self,

{flashrag_dev-0.1.1 → flashrag-dev-0.1.2.dev20241101}/flashrag/refiner/refiner.py

@@ -174,8 +174,8 @@ class ExtractiveRefiner(BaseRefiner):
             
             flatten_batch_sents = sum(batch_sents, [])
             sent_embs = []
-            for s_index in tqdm(range(0, len(flatten_batch_sents), self.mini_batchsize), desc='Sentence encoding..,'):
-                mini_batch_sents = flatten_batch_sents[s_index:s_index+self.mini_batchsize]
+            for s_index in tqdm(range(0, len(flatten_batch_sents), self.mini_batch_size), desc='Sentence encoding..,'):
+                mini_batch_sents = flatten_batch_sents[s_index:s_index+self.mini_batch_size]
                 mini_sent_embs = self.encoder.encode(mini_batch_sents, is_query=False)
                 sent_embs.append(mini_sent_embs)
             sent_embs = np.concatenate(sent_embs, axis=0)

{flashrag_dev-0.1.1 → flashrag-dev-0.1.2.dev20241101}/flashrag/retriever/encoder.py

@@ -1,63 +1,23 @@
 from typing import List
 import torch
 import numpy as np
-from flashrag.retriever.utils import load_model, pooling
-
-
-def parse_query(model_name, query_list, is_query=True):
-    """
-    processing query for different encoders
-    """
-
-    def is_zh(str):
-        import unicodedata
-
-        zh_char = 0
-        for c in str:
-            try:
-                if "CJK" in unicodedata.name(c):
-                    zh_char += 1
-            except:
-                continue
-        if zh_char / len(str) > 0.2:
-            return True
-        else:
-            return False
-
-    if isinstance(query_list, str):
-        query_list = [query_list]
-
-    if "e5" in model_name.lower():
-        if is_query:
-            query_list = [f"query: {query}" for query in query_list]
-        else:
-            query_list = [f"passage: {query}" for query in query_list]
-
-    if "bge" in model_name.lower():
-        if is_query:
-            if is_zh(query_list[0]):
-                query_list = [f"为这个句子生成表示以用于检索相关文章：{query}" for query in query_list]
-            else:
-                query_list = [
-                    f"Represent this sentence for searching relevant passages: {query}" for query in query_list
-                ]
-
-    return query_list
+from flashrag.retriever.utils import load_model, pooling, parse_query
 
 
 class Encoder:
-    def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16):
+    def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16, instruction):
         self.model_name = model_name
         self.model_path = model_path
         self.pooling_method = pooling_method
         self.max_length = max_length
         self.use_fp16 = use_fp16
+        self.instruction = instruction
 
         self.model, self.tokenizer = load_model(model_path=model_path, use_fp16=use_fp16)
 
     @torch.inference_mode(mode=True)
     def encode(self, query_list: List[str], is_query=True) -> np.ndarray:
-        query_list = parse_query(self.model_name, query_list, is_query)
+        query_list = parse_query(self.model_name, query_list, self.instruction)
 
         inputs = self.tokenizer(
             query_list, max_length=self.max_length, padding=True, truncation=True, return_tensors="pt"
@@ -84,7 +44,7 @@ class Encoder:
 
 
 class STEncoder:
-    def __init__(self, model_name, model_path, max_length, use_fp16):
+    def __init__(self, model_name, model_path, max_length, use_fp16, instruction):
         import torch
         from sentence_transformers import SentenceTransformer
 
@@ -92,16 +52,16 @@ class STEncoder:
         self.model_path = model_path
         self.max_length = max_length
         self.use_fp16 = use_fp16
-
+        self.instruction = instruction
         self.model = SentenceTransformer(
             model_path, model_kwargs={"torch_dtype": torch.float16 if use_fp16 else torch.float}
         )
 
     @torch.inference_mode(mode=True)
     def encode(self, query_list: List[str], batch_size=64, is_query=True) -> np.ndarray:
-        query_list = parse_query(self.model_name, query_list, is_query)
+        query_list = parse_query(self.model_name, query_list, self.instruction)
         query_emb = self.model.encode(
-            query_list, batch_size=batch_size, convert_to_numpy=True, normalize_embeddings=True
+            query_list, batch_size=batch_size, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=True
         )
         query_emb = query_emb.astype(np.float32, order="C")
 
@@ -109,10 +69,10 @@ class STEncoder:
 
     @torch.inference_mode(mode=True)
     def multi_gpu_encode(self, query_list: List[str], is_query=True, batch_size=None) -> np.ndarray:
-        query_list = parse_query(self.model_name, query_list, is_query)
+        query_list = parse_query(self.model_name, query_list, self.instruction)
         pool = self.model.start_multi_process_pool()
         query_emb = self.model.encode_multi_process(
-            query_list, pool, convert_to_numpy=True, normalize_embeddings=True, batch_size=batch_size
+            query_list, pool, convert_to_numpy=True, normalize_embeddings=True, batch_size=batch_size,show_progress_bar=True
         )
         self.model.stop_multi_process_pool(pool)
         query_emb.astype(np.float32, order="C")

{flashrag_dev-0.1.1 → flashrag-dev-0.1.2.dev20241101}/flashrag/retriever/index_builder.py

@@ -10,7 +10,7 @@ import argparse
 import datasets
 import torch
 from tqdm import tqdm
-from flashrag.retriever.utils import load_model, load_corpus, pooling
+from flashrag.retriever.utils import load_model, load_corpus, pooling, set_default_instruction
 
 
 class Index_Builder:
@@ -25,7 +25,8 @@ class Index_Builder:
         max_length,
         batch_size,
         use_fp16,
-        pooling_method,
+        pooling_method=None,
+        instruction=None,
         faiss_type=None,
         embedding_path=None,
         save_embedding=False,
@@ -41,7 +42,7 @@ class Index_Builder:
         self.max_length = max_length
         self.batch_size = batch_size
         self.use_fp16 = use_fp16
-        self.pooling_method = pooling_method
+        self.instruction = instruction
         self.faiss_type = faiss_type if faiss_type is not None else "Flat"
         self.embedding_path = embedding_path
         self.save_embedding = save_embedding
@@ -49,6 +50,43 @@ class Index_Builder:
         self.use_sentence_transformer = use_sentence_transformer
         self.bm25_backend = bm25_backend
 
+        # set instruction for encode
+        if self.instruction is not None:
+            self.instruction = self.instruction.strip() + " "
+            print("Set instruction for encoding:", self.instruction)
+        else:
+            self.instruction = set_default_instruction(self.retrieval_method, is_query=False)
+            if self.instruction == "":
+                warnings.warn("Instruction is not set!")
+            else:
+                warnings.warn(f"Instruction is set to default: {self.instruction}")
+
+        #. config pooling method
+        if pooling_method is None:
+            try:
+                # read pooling method from 1_Pooling/config.json
+                pooling_config = json.load(open(os.path.join(self.model_path, "1_Pooling/config.json")))
+                for k, v in pooling_config.items():
+                    if k.startswith("pooling_mode") and v == True:
+                        pooling_method = k.split("pooling_mode_")[-1]
+                        if pooling_method == 'mean_tokens':
+                            pooling_method = 'mean'
+                        elif pooling_method == 'cls_token':
+                            pooling_method = 'cls'
+                        else:
+                            # raise warning: not implemented pooling method
+                            warnings.warn(f"Pooling method {pooling_method} is not implemented.", UserWarning)
+                            pooling_method = 'mean'
+                        break
+            except:
+                print(f"Pooling method not found in {self.model_path}, use default pooling method (mean).")
+                # use default pooling method
+                pooling_method = 'mean'
+        else:
+            if pooling_method not in ['mean', 'cls', 'pooler']:
+                raise ValueError(f"Invalid pooling method {pooling_method}.")
+        self.pooling_method = pooling_method
+
         self.gpu_num = torch.cuda.device_count()
         # prepare save dir
         print(self.save_dir)
@@ -162,8 +200,7 @@ class Index_Builder:
             self.batch_size = self.batch_size * self.gpu_num
 
         sentence_list = [item["contents"] for item in self.corpus]
-        if self.retrieval_method == "e5":
-            sentence_list = [f"passage: {doc}" for doc in sentence_list]
+        sentence_list = [f"{self.instruction}{doc}" for doc in sentence_list]
         all_embeddings = self.encoder.encode(sentence_list, batch_size=self.batch_size)
 
         return all_embeddings
@@ -175,12 +212,9 @@ class Index_Builder:
             self.batch_size = self.batch_size * self.gpu_num
 
         all_embeddings = []
-
         for start_idx in tqdm(range(0, len(self.corpus), self.batch_size), desc="Inference Embeddings:"):
             batch_data = self.corpus[start_idx : start_idx + self.batch_size]["contents"]
-
-            if self.retrieval_method == "e5":
-                batch_data = [f"passage: {doc}" for doc in batch_data]
+            batch_data = [f"{self.instruction}{doc}" for doc in batch_data]
 
             inputs = self.tokenizer(
                 batch_data,
@@ -273,9 +307,6 @@ class Index_Builder:
         print("Finish!")
 
 
-MODEL2POOLING = {"e5": "mean", "bge": "cls", "contriever": "mean", "jina": "mean"}
-
-
 def main():
     parser = argparse.ArgumentParser(description="Creating index.")
 
@@ -290,6 +321,7 @@ def main():
     parser.add_argument("--batch_size", type=int, default=512)
     parser.add_argument("--use_fp16", default=False, action="store_true")
     parser.add_argument("--pooling_method", type=str, default=None)
+    parser.add_argument("--instruction", type=str, default=None)
     parser.add_argument("--faiss_type", default=None, type=str)
     parser.add_argument("--embedding_path", default=None, type=str)
     parser.add_argument("--save_embedding", action="store_true", default=False)
@@ -299,18 +331,6 @@ def main():
 
     args = parser.parse_args()
 
-    if args.pooling_method is None:
-        pooling_method = "mean"
-        for k, v in MODEL2POOLING.items():
-            if k in args.retrieval_method.lower():
-                pooling_method = v
-                break
-    else:
-        if args.pooling_method not in ["mean", "cls", "pooler"]:
-            raise NotImplementedError
-        else:
-            pooling_method = args.pooling_method
-
     index_builder = Index_Builder(
         retrieval_method=args.retrieval_method,
         model_path=args.model_path,
@@ -319,7 +339,8 @@ def main():
         max_length=args.max_length,
         batch_size=args.batch_size,
         use_fp16=args.use_fp16,
-        pooling_method=pooling_method,
+        pooling_method=args.pooling_method,
+        instruction=args.instruction,
         faiss_type=args.faiss_type,
         embedding_path=args.embedding_path,
         save_embedding=args.save_embedding,

{flashrag_dev-0.1.1 → flashrag-dev-0.1.2.dev20241101}/flashrag/retriever/retriever.py

@@ -38,7 +38,7 @@ def cache_manager(func):
                         warnings.warn(f"The number of cached retrieval results is less than topk ({num})")
                     cache_res = cache_res[:num]
                     # separate the doc score
-                    doc_scores = [item.pop("score") for item in cache_res]
+                    doc_scores = [item["score"] for item in cache_res]
                     cache_results.append((cache_res, doc_scores))
                 else:
                     cache_results.append(None)
@@ -194,10 +194,10 @@ class BM25Retriever(BaseRetriever):
             import Stemmer
             import bm25s
 
-            stemmer = Stemmer.Stemmer('english')
+            self.stemmer = Stemmer.Stemmer('english')
             self.searcher = bm25s.BM25.load(self.index_path, mmap=True, load_corpus=True)
             self.searcher.backend = 'numba'
-            self.tokenizer = bm25s.tokenization.Tokenizer(stemmer=stemmer)
+            
         else:
             assert False, 'Invalid bm25 backend!'
 
@@ -235,7 +235,8 @@ class BM25Retriever(BaseRetriever):
             else:
                 results = load_docs(self.corpus, [hit.docid for hit in hits])
         elif self.backend == 'bm25s':
-            query_tokens = self.tokenizer.tokenize([query])
+            import bm25s 
+            query_tokens = bm25s.tokenize([query], stemmer=self.stemmer)
             results, scores = self.searcher.retrieve(query_tokens, k=num)
             results = results[0]
             scores = scores[0]
@@ -257,7 +258,8 @@ class BM25Retriever(BaseRetriever):
                 results.append(item_result)
                 scores.append(item_score)
         elif self.backend == 'bm25s':
-            query_tokens = self.tokenizer.tokenize(query_list)
+            import bm25s
+            query_tokens = bm25s.tokenize(query_list, stemmer=self.stemmer)
             results, scores = self.searcher.retrieve(query_tokens, k=num)
         else:
             assert False, 'Invalid bm25 backend!'
@@ -273,6 +275,8 @@ class DenseRetriever(BaseRetriever):
 
     def __init__(self, config: dict):
         super().__init__(config)
+        if not os.path.exists(self.index_path):
+            raise Warning(f"Index file {self.index_path} does not exist!")
         self.index = faiss.read_index(self.index_path)
         if config["faiss_gpu"]:
             co = faiss.GpuMultipleClonerOptions()
@@ -281,6 +285,9 @@ class DenseRetriever(BaseRetriever):
             self.index = faiss.index_cpu_to_all_gpus(self.index, co=co)
 
         self.corpus = load_corpus(self.corpus_path)
+        self.topk = config["retrieval_topk"]
+        self.batch_size = config["retrieval_batch_size"]
+        self.instruction = config["instruction"]
 
         if config["use_sentence_transformer"]:
             self.encoder = STEncoder(
@@ -288,6 +295,7 @@ class DenseRetriever(BaseRetriever):
                 model_path=config["retrieval_model_path"],
                 max_length=config["retrieval_query_max_length"],
                 use_fp16=config["retrieval_use_fp16"],
+                instruction=self.instruction,
             )
         else:
             self.encoder = Encoder(
@@ -296,9 +304,9 @@ class DenseRetriever(BaseRetriever):
                 pooling_method=config["retrieval_pooling_method"],
                 max_length=config["retrieval_query_max_length"],
                 use_fp16=config["retrieval_use_fp16"],
+                instruction=self.instruction,
             )
-        self.topk = config["retrieval_topk"]
-        self.batch_size = self.config["retrieval_batch_size"]
+        
 
     def _search(self, query: str, num: int = None, return_score=False):
         if num is None:

{flashrag_dev-0.1.1 → flashrag-dev-0.1.2.dev20241101}/flashrag/retriever/utils.py

@@ -26,6 +26,55 @@ def pooling(pooler_output, last_hidden_state, attention_mask=None, pooling_metho
     else:
         raise NotImplementedError("Pooling method not implemented!")
 
+def set_default_instruction(model_name, is_query=True, is_zh=False):
+    instruction = ""
+    if "e5" in model_name.lower():
+        if is_query:
+            instruction = "query: "
+        else:
+            instruction = "passage: "
+
+    if "bge" in model_name.lower():
+        if is_query:
+            if 'zh' in model_name.lower() or is_zh:
+                instruction = "为这个句子生成表示以用于检索相关文章："
+            else:
+                instruction = "Represent this sentence for searching relevant passages: "
+
+    return instruction
+
+
+def parse_query(model_name, query_list, instruction=None):
+    """
+    processing query for different encoders
+    """
+
+    def is_zh(str):
+        import unicodedata
+
+        zh_char = 0
+        for c in str:
+            try:
+                if "CJK" in unicodedata.name(c):
+                    zh_char += 1
+            except:
+                continue
+        if zh_char / len(str) > 0.2:
+            return True
+        else:
+            return False
+
+    if isinstance(query_list, str):
+        query_list = [query_list]
+        
+    if instruction is not None:
+        instruction = instruction.strip() + " "
+    else:
+        instruction = set_default_instruction(model_name, is_query=True, is_zh=is_zh(query_list[0]))
+
+    query_list = [instruction + query for query in query_list]
+
+    return query_list
 
 def load_corpus(corpus_path: str):
     corpus = datasets.load_dataset("json", data_files=corpus_path, split="train")

{flashrag_dev-0.1.1 → flashrag-dev-0.1.2.dev20241101}/flashrag/utils/utils.py

@@ -101,7 +101,9 @@ def get_refiner(config, retriever=None, generator=None):
     try:
         model_config = AutoConfig.from_pretrained(refiner_path)
         arch = model_config.architectures[0].lower()
-    except:
+        print(arch)
+    except Exception as e:
+        print("Warning", e)
         model_config, arch = "", ""
 
     if "recomp" in refiner_name or "bert" in arch:

flashrag-dev-0.1.2.dev20241101/flashrag/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.2dev20241101"

{flashrag_dev-0.1.1 → flashrag-dev-0.1.2.dev20241101/flashrag_dev.egg-info}/PKG-INFO

@@ -1,75 +1,18 @@
 Metadata-Version: 2.1
 Name: flashrag-dev
-Version: 0.1.1
+Version: 0.1.2.dev20241101
 Summary: A library for efficient Retrieval-Augmented Generation research
 Home-page: https://github.com/RUC-NLPIR/FlashRAG
 Author: Jiajie Jin, Yutao Zhu, Chenghao Zhang, Xinyu Yang, Zhicheng Dou
 Author-email: jinjiajie@ruc.edu.cn
 License: MIT License
-Requires-Python: >=3.8
+Requires-Python: >=3.9
 Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: datasets
-Requires-Dist: base58
-Requires-Dist: nltk
-Requires-Dist: numpy
-Requires-Dist: openai
-Requires-Dist: peft
-Requires-Dist: PyYAML
-Requires-Dist: rank_bm25
-Requires-Dist: rouge
-Requires-Dist: spacy>=3.6.1
-Requires-Dist: tiktoken
-Requires-Dist: torch
-Requires-Dist: tqdm
-Requires-Dist: transformers>=4.40.0
-Requires-Dist: bm25s[full]
-Requires-Dist: fschat
-Requires-Dist: streamlit
 Provides-Extra: core
-Requires-Dist: datasets; extra == "core"
-Requires-Dist: base58; extra == "core"
-Requires-Dist: nltk; extra == "core"
-Requires-Dist: numpy; extra == "core"
-Requires-Dist: openai; extra == "core"
-Requires-Dist: peft; extra == "core"
-Requires-Dist: PyYAML; extra == "core"
-Requires-Dist: rank_bm25; extra == "core"
-Requires-Dist: rouge; extra == "core"
-Requires-Dist: spacy>=3.6.1; extra == "core"
-Requires-Dist: tiktoken; extra == "core"
-Requires-Dist: torch; extra == "core"
-Requires-Dist: tqdm; extra == "core"
-Requires-Dist: transformers>=4.40.0; extra == "core"
-Requires-Dist: bm25s[full]; extra == "core"
-Requires-Dist: fschat; extra == "core"
-Requires-Dist: streamlit; extra == "core"
 Provides-Extra: retriever
-Requires-Dist: pyserini; extra == "retriever"
-Requires-Dist: sentence-transformers>=3.0.1; extra == "retriever"
 Provides-Extra: generator
-Requires-Dist: vllm>=0.4.1; extra == "generator"
 Provides-Extra: full
-Requires-Dist: datasets; extra == "full"
-Requires-Dist: base58; extra == "full"
-Requires-Dist: nltk; extra == "full"
-Requires-Dist: numpy; extra == "full"
-Requires-Dist: openai; extra == "full"
-Requires-Dist: peft; extra == "full"
-Requires-Dist: PyYAML; extra == "full"
-Requires-Dist: rank_bm25; extra == "full"
-Requires-Dist: rouge; extra == "full"
-Requires-Dist: spacy>=3.6.1; extra == "full"
-Requires-Dist: tiktoken; extra == "full"
-Requires-Dist: torch; extra == "full"
-Requires-Dist: tqdm; extra == "full"
-Requires-Dist: transformers>=4.40.0; extra == "full"
-Requires-Dist: bm25s[full]; extra == "full"
-Requires-Dist: fschat; extra == "full"
-Requires-Dist: streamlit; extra == "full"
-Requires-Dist: pyserini; extra == "full"
-Requires-Dist: sentence-transformers>=3.0.1; extra == "full"
-Requires-Dist: vllm>=0.4.1; extra == "full"
+License-File: LICENSE
 
 # <div align="center">⚡FlashRAG: A Python Toolkit for Efficient RAG Research<div>
 
@@ -94,7 +37,7 @@ Requires-Dist: vllm>=0.4.1; extra == "full"
 </p>
 
 </h4>
-FlashRAG is a Python toolkit for the reproduction and development of Retrieval Augmented Generation (RAG) research. Our toolkit includes 32 pre-processed benchmark RAG datasets and 15 state-of-the-art RAG algorithms. 
+FlashRAG is a Python toolkit for the reproduction and development of Retrieval Augmented Generation (RAG) research. Our toolkit includes 33 pre-processed benchmark RAG datasets and 15 state-of-the-art RAG algorithms. 
 
 <p align="center">
 <img src="asset/framework.jpg">
@@ -110,7 +53,7 @@ With FlashRAG and provided resources, you can effortlessly reproduce existing SO
 
 - **Extensive and Customizable Framework**: Includes essential components for RAG scenarios such as retrievers, rerankers, generators, and compressors, allowing for flexible assembly of complex pipelines.
 
-- **Comprehensive Benchmark Datasets**: A collection of 32 pre-processed RAG benchmark datasets to test and validate RAG models' performances.
+- **Comprehensive Benchmark Datasets**: A collection of 33 pre-processed RAG benchmark datasets to test and validate RAG models' performances.
 
 - **Pre-implemented Advanced RAG Algorithms**: Features 15 advancing RAG algorithms with reported results, based on our framework. Easily reproducing results under different settings.
 
@@ -123,7 +66,6 @@ With FlashRAG and provided resources, you can effortlessly reproduce existing SO
 FlashRAG is still under development and there are many issues and room for improvement. We will continue to update. And we also sincerely welcome contributions on this open-source toolkit.
 
 - [x] Support OpenAI models
-- [ ] Support Claude and Gemini models
 - [x] Provdide instructions for each component
 - [x] Integrate sentence Transformers
 - [ ] Inlcude more RAG approaches
@@ -132,6 +74,11 @@ FlashRAG is still under development and there are many issues and room for impro
 
 
 ## :page_with_curl: Changelog
+[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.
+
+[24/10/13] A new in-domain dataset and corpus - [DomainRAG](https://arxiv.org/pdf/2406.05654) have been added to the dataset. The dataset is based on the internal enrollment data of Renmin University of China, covering seven types of tasks, which can be used for conducting domain-specific RAG testing.
+
+[24/09/24] We have released a version based on the MindSpore framework that supports Chinese hardware platforms. Please refer to [FlashRAG MindSpore](https://github.com/RUC-NLPIR/FlashRAG-MindSpore) for details.
 
 [24/09/18] Due to the complexity and limitations of installing Pyserini in certain environments, we have introduced a lightweight `BM25s` package as an alternative (faster and easier to use). The retriever based on Pyserini will be deprecated in future versions. To use retriever with `bm25s`, just set `bm25_backend` to `bm25s` in config.
 
@@ -165,7 +112,7 @@ FlashRAG is still under development and there are many issues and room for impro
 
 To get started with FlashRAG, you can simply install it with pip:
 ```base
-pip install flashrag[core]
+pip install flashrag-dev[core]
 ```
 
 Or you can clone it from Github and install (requires Python 3.9+): 
@@ -554,6 +501,7 @@ Below is the list of datasets along with the corresponding sample sizes:
 | multi-hop QA              | 2WikiMultiHopQA | wiki             | 15,000    | 12,576  | /      |
 | multi-hop QA              | Musique         | wiki             | 19,938    | 2,417   | /      |
 | multi-hop QA              | Bamboogle       | wiki             | /         | /       | 125    |
+| multi-hop QA              | StrategyQA      | wiki             | 2290      | /       | /
 | Long-form QA              | ASQA            | wiki             | 4,353     | 948     | /      |
 | Long-form QA              | ELI5            | Reddit           | 272,634   | 1,507   | /      |
 | Open-Domain Summarization | WikiASP         | wiki             | 300,636   | 37,046  | 37,368 |
@@ -568,6 +516,7 @@ Below is the list of datasets along with the corresponding sample sizes:
 | Entity Linking            | WNED            | Wiki             | /         | 8,995   | /      |
 | Slot Filling              | T-REx           | DBPedia          | 2,284,168 | 5,000   | /      |
 | Slot Filling              | Zero-shot RE    | wiki             | 147,909   | 3,724   | /      |
+| In-domain QA| DomainRAG | Web pages of RUC| / | / | 485|
 
 ### Document Corpus
 

{flashrag_dev-0.1.1 → flashrag-dev-0.1.2.dev20241101}/flashrag_dev.egg-info/SOURCES.txt

@@ -3,7 +3,9 @@ README.md
 pyproject.toml
 setup.py
 flashrag/__init__.py
+flashrag/version.py
 flashrag/config/__init__.py
+flashrag/config/basic_config.yaml
 flashrag/config/config.py
 flashrag/dataset/__init__.py
 flashrag/dataset/dataset.py

{flashrag_dev-0.1.1 → flashrag-dev-0.1.2.dev20241101}/flashrag_dev.egg-info/requires.txt

@@ -55,10 +55,10 @@ fschat
 streamlit
 pyserini
 sentence-transformers>=3.0.1
-vllm>=0.4.1
+vllm==0.5.5
 
 [generator]
-vllm>=0.4.1
+vllm==0.5.5
 
 [retriever]
 pyserini

{flashrag_dev-0.1.1 → flashrag-dev-0.1.2.dev20241101}/setup.py

@@ -4,17 +4,20 @@ with open("requirements.txt",encoding='utf-8') as fp:
     requirements = fp.read().splitlines()
 with open("README.md", "r",encoding='utf-8') as fh:
     long_description = fh.read()
+version = {}
+with open("flashrag/version.py", encoding="utf8") as fp:
+    exec(fp.read(), version)
 
 extras_require = {
     'core': requirements,
     'retriever': ['pyserini', 'sentence-transformers>=3.0.1'],
-    'generator': ['vllm>=0.4.1']
+    'generator': ['vllm==0.5.5']
 }
 extras_require['full'] = sum(extras_require.values(), [])
 
 setup(
     name="flashrag-dev",
-    version="0.1.1",
+    version=version['__version__'],
     packages=find_packages(),
     url="https://github.com/RUC-NLPIR/FlashRAG",
     license="MIT License",
@@ -23,9 +26,9 @@ setup(
     description="A library for efficient Retrieval-Augmented Generation research",
     long_description=long_description,
     long_description_content_type="text/markdown",
+    package_data={"": ["*.yaml"]},
     include_package_data=True,
-    package_data={"flashrag-dev": ["/config/basic_config.yaml"]},
     install_requires=extras_require['core'],
     extras_require=extras_require,
-    python_requires=">=3.8",
+    python_requires=">=3.9",
 )