flashrag-dev 0.1.4.dev20250715__tar.gz → 0.1.4.dev20250805__tar.gz

{flashrag_dev-0.1.4.dev20250715 → flashrag_dev-0.1.4.dev20250805}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag_dev
-Version: 0.1.4.dev20250715
+Version: 0.1.4.dev20250805
 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
@@ -278,6 +278,32 @@ python -m flashrag.retriever.index_builder \
   --save_dir indexes/ 
 ```
 
+### For Sparse Neural Retrieval Methods (SPLADE)
+
+##### Install Seismic Index:
+```bash
+curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh # Install Rust for compiling
+pip install pyseismic-lsr # Install Seismic
+```
+
+##### Then build the index with Seismic:
+```bash
+python -m flashrag.retriever.index_builder \ # builder
+        --retrieval_method splade \ # Model name to trigger seismic index (splade only available)
+        --model_path retriever/splade-v3 \ # Local path or repository path are both supported.
+        --corpus_embedded_path data/ms_marco/ms_marco_embedded_corpus.jsonl \  # Use cached embedded corpus if corpus is already available in seismic expected format
+        --corpus_path data/ms_marco/ms_marco_corpus.jsonl \ # Corpus path in format {id, contents} jsonl file to be embedded if not already built
+        --save_dir indexes/ \ # save index directory
+        --use_fp16 \ # tell to use fp16 for splade model
+        --max_length 512 \ # max tokens for each document
+        --batch_size 4 \ # batch size for splade model (4-5 seems the best size for Tesla T4 16GB)
+        --n_postings 1000 \ # seismic number of posting lists
+        --centroid_fraction 0.2 \ # seismic centroids
+        --min_cluster_size 2 \ # seismic min cluster
+        --summary_energy 0.4 \ # seismic energy
+        --batched_indexing 10000000 # seismic batch
+        --nknn 32 # Optional parameter. Tell to seismic to use also knn graph. if not present seismic will work without knn graph
+```
 ### Using the ready-made pipeline
 
 You can use the pipeline class we have already built (as shown in [<u>pipelines</u>](#pipelines)) to implement the RAG process inside. In this case, you just need to configure the config and load the corresponding pipeline.

flashrag_dev-0.1.4.dev20250715/flashrag_dev.egg-info/PKG-INFO → flashrag_dev-0.1.4.dev20250805/README.md

@@ -1,20 +1,3 @@
-Metadata-Version: 2.1
-Name: flashrag-dev
-Version: 0.1.4.dev20250715
-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.9
-Description-Content-Type: text/markdown
-Provides-Extra: core
-Provides-Extra: retriever
-Provides-Extra: generator
-Provides-Extra: multimodal
-Provides-Extra: full
-License-File: LICENSE
-
 # <div align="center">⚡FlashRAG: A Python Toolkit for Efficient RAG Research<div>
 \[ English | [中文](README_zh.md) \]
 <div align="center">
@@ -278,6 +261,32 @@ python -m flashrag.retriever.index_builder \
   --save_dir indexes/ 
 ```
 
+### For Sparse Neural Retrieval Methods (SPLADE)
+
+##### Install Seismic Index:
+```bash
+curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh # Install Rust for compiling
+pip install pyseismic-lsr # Install Seismic
+```
+
+##### Then build the index with Seismic:
+```bash
+python -m flashrag.retriever.index_builder \ # builder
+        --retrieval_method splade \ # Model name to trigger seismic index (splade only available)
+        --model_path retriever/splade-v3 \ # Local path or repository path are both supported.
+        --corpus_embedded_path data/ms_marco/ms_marco_embedded_corpus.jsonl \  # Use cached embedded corpus if corpus is already available in seismic expected format
+        --corpus_path data/ms_marco/ms_marco_corpus.jsonl \ # Corpus path in format {id, contents} jsonl file to be embedded if not already built
+        --save_dir indexes/ \ # save index directory
+        --use_fp16 \ # tell to use fp16 for splade model
+        --max_length 512 \ # max tokens for each document
+        --batch_size 4 \ # batch size for splade model (4-5 seems the best size for Tesla T4 16GB)
+        --n_postings 1000 \ # seismic number of posting lists
+        --centroid_fraction 0.2 \ # seismic centroids
+        --min_cluster_size 2 \ # seismic min cluster
+        --summary_energy 0.4 \ # seismic energy
+        --batched_indexing 10000000 # seismic batch
+        --nknn 32 # Optional parameter. Tell to seismic to use also knn graph. if not present seismic will work without knn graph
+```
 ### Using the ready-made pipeline
 
 You can use the pipeline class we have already built (as shown in [<u>pipelines</u>](#pipelines)) to implement the RAG process inside. In this case, you just need to configure the config and load the corresponding pipeline.

{flashrag_dev-0.1.4.dev20250715 → flashrag_dev-0.1.4.dev20250805}/flashrag/config/basic_config.yaml

@@ -68,7 +68,9 @@ bm25_backend: bm25s # pyserini, bm25s
 use_sentence_transformer: False
 silent_retrieval: True # whether to silent the retrieval process
 
-
+seismic_query_cut: 10 # parameters for seismic. See seismic paper for full details
+seismic_heap_factor: 0.8 # parameters for seismic. See seismic paper for full details
+# -------------------------------------------------Reranker Settings------------------------------------------------#
 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`
@@ -117,6 +119,7 @@ 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

{flashrag_dev-0.1.4.dev20250715 → flashrag_dev-0.1.4.dev20250805}/flashrag/generator/generator.py

@@ -12,6 +12,7 @@ from transformers import (
     AutoConfig,
 )
 from flashrag.generator.utils import resolve_max_tokens
+from flashrag.utils import get_device
 
 
 class BaseGenerator:
@@ -553,13 +554,13 @@ class FastChatGenerator(HFCausalLMGenerator):
             max_gpu_memory = None
             import torch
             self.gpu_num = torch.cuda.device_count()
-            if self.gpu_num != 1:
+            if self.gpu_num > 1:
                 available_gpu_memory = get_gpu_memory()
                 max_gpu_memory = str(int(min(available_gpu_memory) * gpu_memory_utilization)) + "GiB"
 
             model, tokenizer = load_model(
                 self.model_path,
-                device="cuda",
+                device=get_device(),
                 num_gpus=self.gpu_num,
                 max_gpu_memory=max_gpu_memory,
                 load_8bit=False,

flashrag_dev-0.1.4.dev20250805/flashrag/pipeline/ReaRAG_utils.py

@@ -0,0 +1,150 @@
+import re
+from termcolor import colored
+
+class AgentUtilsBase():
+    def __init__(self):
+        pass
+
+    def truncate(self, tokenizer, tokenized_prompt, model_max_length):
+        half = int(model_max_length/2)
+        prompt = tokenizer.decode(tokenized_prompt[:half], skip_special_tokens=False) + \
+                tokenizer.decode(tokenized_prompt[-half:], skip_special_tokens=False)
+        
+        return prompt
+
+    def preprocess_query(self, query):
+        if "'" in query and '"' in query:
+            query = query.replace("'", "\\'").replace('"', '\\"')
+        return query
+
+    def extract_code(self, text: str) -> str:
+        triple_match = re.search(r'```[^\n]*\n(.+?)```', text, re.DOTALL)
+        single_match = re.search(r'`([^`]*)`', text, re.DOTALL)
+        if triple_match:
+            return triple_match.group(1)
+        elif single_match:
+            return single_match.group(1)
+        return text
+    
+    def postprocess_agent_response(self, response):
+        """
+        Implement the postprocess_agent_response function
+        """
+        raise NotImplementedError("postprocess_agent_response function must be implemented in a subclass")
+
+class AgentUtils(AgentUtilsBase):
+    def __init__(self):
+        super().__init__()
+
+    def parse_reasoning_steps(self, text: str):
+        """
+        Parse a string containing Thought/Action/Observation steps (including multi-line)
+        and return a list of dictionaries of the form:
+        
+        [
+            {
+                "1": {
+                    "Thought": "...",
+                    "Action": "...",  # Only content inside backticks (if present)
+                    "Observation": "..."
+                }
+            },
+            {
+                "2": {
+                    "Thought": "...",
+                    "Action": "...",
+                    "Observation": "..."
+                }
+            },
+            ...
+        ]
+        """
+        # Regex pattern to match lines that start with "Thought X:", "Action X:", or "Observation X:".
+        pattern = re.compile(r'^(Thought|Action|Observation)\s+(\d+):', re.MULTILINE)
+
+        # This dictionary will accumulate:
+        # data_dict[step_number] = {"Thought": ..., "Action": ..., "Observation": ...}
+        data_dict = {}
+
+        # We'll track the current label (Thought/Action/Observation) and step number
+        current_label = None
+        current_step = None
+        last_pos = 0
+
+        # Find all pattern occurrences in the text
+        matches = list(pattern.finditer(text))
+
+        for match in matches:
+            # If we already have a label in progress, we can record its content
+            if current_label is not None:
+                # Slice the text from the last match's end to the start of this new match
+                content = text[last_pos:match.start()].strip()
+                # Store that content in data_dict
+                data_dict[current_step][current_label] = content
+
+            # Extract the new label and step
+            label = match.group(1)       # "Thought", "Action", or "Observation"
+            step = match.group(2)        # e.g. "1", "2", "3"
+
+            # Ensure a dict for this step
+            if step not in data_dict:
+                data_dict[step] = {"Thought": None, "Action": None, "Observation": None}
+
+            # Update current label/step
+            current_label = label
+            current_step = step
+            # We'll slice from here next time
+            last_pos = match.end()
+
+        # Handle the final block after the last match
+        if current_label is not None:
+            content = text[last_pos:].strip()
+            data_dict[current_step][current_label] = content
+
+        # Post-process:
+        #  - For each step, extract only the text inside triple backticks for "Action".
+        for step_number in data_dict:
+            action_text = data_dict[step_number]["Action"]
+            if action_text:
+                # Extract content inside triple backticks
+                data_dict[step_number]["Action"] = self.extract_code(action_text)
+
+        # Convert our dictionary to the desired list-of-dicts structure
+        structured_data = []
+        for step_number in sorted(data_dict.keys(), key=lambda x: int(x)):
+            structured_data.append({step_number: data_dict[step_number]})
+
+        return structured_data
+    
+    def postprocess_agent_response(self, response):
+        """
+        Extract Thought and Action, then extract the dict from Action.
+
+        Return:
+        - thought: str
+        - action: dict
+        - is_valid: bool
+        """
+        parsed_codes = self.parse_reasoning_steps(response)
+
+        thoughts = []
+        actions = []
+        for steps in parsed_codes:
+            for step_idx, step in steps.items():
+                thought = f"Thought {step_idx}: {step['Thought']}"
+                action = eval(self.extract_code(step['Action']))
+
+                assert "function" in action, f"Action does not contain 'function' key: {action}"
+                assert "parameters" in action, f"Action does not contain 'parameters' key: {action}"
+
+                thoughts.append(thought)
+                actions.append(action)
+
+        return thoughts, actions
+    
+def print_code(codes):
+    for idx, step in enumerate(codes):
+        print(f"{colored(step['thought'], 'blue')}")
+        print(colored(f"Action {idx+1}:\n```\n{step['action']}\n```", 'red'))
+        print(colored(f"Observation {idx+1}: {step['observation']}", 'yellow'))
+        # print(f"{'-'*60}")

{flashrag_dev-0.1.4.dev20250715 → flashrag_dev-0.1.4.dev20250805}/flashrag/pipeline/__init__.py

@@ -2,4 +2,5 @@ from flashrag.pipeline.mm_pipeline import *
 from flashrag.pipeline.pipeline import *
 from flashrag.pipeline.branching_pipeline import REPLUGPipeline, SuRePipeline
 from flashrag.pipeline.active_pipeline import IterativePipeline, SelfRAGPipeline, FLAREPipeline, SelfAskPipeline, IRCOTPipeline, RQRAGPipeline
-from flashrag.pipeline.reasoning_pipeline import ReasoningPipeline
+from flashrag.pipeline.reasoning_pipeline import *
+from flashrag.pipeline.ReaRAG_utils import *