flashrag-dev 0.1.3.dev20241108__tar.gz → 0.1.3.dev20241111__tar.gz

{flashrag-dev-0.1.3.dev20241108/flashrag_dev.egg-info → flashrag-dev-0.1.3.dev20241111}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag-dev
-Version: 0.1.3.dev20241108
+Version: 0.1.3.dev20241111
 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

flashrag-dev-0.1.3.dev20241111/flashrag/dataset/dataset.py

@@ -0,0 +1,180 @@
+import os
+import json
+import random
+import warnings
+from typing import List, Dict, Any, Optional, Generator
+import numpy as np
+
+
+class Item:
+    """A container class used to store and manipulate a sample within a dataset.
+    Information related to this sample during training/inference will be stored in `self.output`.
+    Each attribute of this class can be used like a dict key (also for key in `self.output`).
+    """
+
+    def __init__(self, item_dict: Dict[str, Any]) -> None:
+        self.id: Optional[str] = item_dict.get("id", None)
+        self.question: Optional[str] = item_dict.get("question", None)
+        self.golden_answers: List[str] = item_dict.get("golden_answers", [])
+        self.choices: List[str] = item_dict.get("choices", [])
+        self.metadata: Dict[str, Any] = item_dict.get("metadata", {})
+        self.output: Dict[str, Any] = item_dict.get("output", {})
+        self.data: Dict[str, Any] = item_dict
+
+    def update_output(self, key: str, value: Any) -> None:
+        """Update the output dict and keep a key in self.output can be used as an attribute."""
+        if key in ["id", "question", "golden_answers", "output", "choices"]:
+            raise AttributeError(f"{key} should not be changed")
+        else:
+            self.output[key] = value
+
+    def update_evaluation_score(self, metric_name: str, metric_score: float) -> None:
+        """Update the evaluation score of this sample for a metric."""
+        if "metric_score" not in self.output:
+            self.output["metric_score"] = {}
+        self.output["metric_score"][metric_name] = metric_score
+
+    def __getattr__(self, attr_name: str) -> Any:
+        predefined_attrs = ["id", "question", "golden_answers", "metadata", "output", "choices"]
+        if attr_name in predefined_attrs:
+            return super().__getattribute__(attr_name)
+        else:
+            output = self.output
+            if attr_name in output:
+                return output[attr_name]
+            else:
+                try:
+                    return self.data[attr_name]
+                except AttributeError:
+                    raise AttributeError(f"Attribute `{attr_name}` not found")
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert all information within the data sample into a dict. Information generated
+        during the inference will be saved into output field.
+        """
+        from flashrag.dataset.utils import convert_numpy
+
+        output = {
+            "id": self.id,
+            "question": self.question,
+            "golden_answers": self.golden_answers,
+            "output": convert_numpy(self.output),
+        }
+        if self.metadata:
+            output["metadata"] = self.metadata
+
+        return output
+
+    def __str__(self) -> str:
+        """Return a string representation of the item with its main attributes."""
+        return json.dumps(self.to_dict(), indent=4)
+
+
+class Dataset:
+    """A container class used to store the whole dataset. Inside the class, each data sample will be stored
+    in `Item` class. The properties of the dataset represent the list of attributes corresponding to each item in the dataset.
+    """
+
+    def __init__(
+        self,
+        config: Optional[Dict[str, Any]] = None,
+        dataset_path: Optional[str] = None,
+        data: Optional[List[Dict[str, Any]]] = None,
+        sample_num: Optional[int] = None,
+        random_sample: bool = False,
+    ) -> None:
+        if config is not None:
+            self.config = config
+            dataset_name = config['dataset_name'] if 'dataset_name' in config else 'defalut_dataset'
+        else:
+            self.config = None
+            warnings.warn("dataset_name is not in config, set it as default.")
+            dataset_name = "default_dataset"
+        self.dataset_name = dataset_name
+        self.dataset_path = dataset_path
+
+        self.sample_num = sample_num
+        self.random_sample = random_sample
+
+        if data is None:
+            self.data = self._load_data(self.dataset_name, self.dataset_path)
+        else:
+            print("Load data from provided data")
+            assert isinstance(data[0], dict)
+            self.data = [Item(item_dict) for item_dict in data]
+
+    def _load_data(self, dataset_name: str, dataset_path: str) -> List[Item]:
+        """Load data from the provided dataset_path or directly download the file(TODO)."""
+        if not os.path.exists(dataset_path):
+            # TODO: auto download: self._download(self.dataset_name, dataset_path)
+            raise FileNotFoundError(f"Dataset file {dataset_path} not found.")
+
+        data = []
+        with open(dataset_path, "r", encoding="utf-8") as f:
+            for line in f:
+                item_dict = json.loads(line)
+                item = Item(item_dict)
+                data.append(item)
+        if self.sample_num is not None:
+            if self.random_sample:
+                print(f"Random sample {self.sample_num} items in test set.")
+                data = random.sample(data, self.sample_num)
+            else:
+                data = data[: self.sample_num]
+
+        return data
+
+    def update_output(self, key: str, value_list: List[Any]) -> None:
+        """Update the overall output field for each sample in the dataset."""
+        assert len(self.data) == len(value_list)
+        for item, value in zip(self.data, value_list):
+            item.update_output(key, value)
+
+    @property
+    def question(self) -> List[Optional[str]]:
+        return [item.question for item in self.data]
+
+    @property
+    def golden_answers(self) -> List[List[str]]:
+        return [item.golden_answers for item in self.data]
+
+    @property
+    def id(self) -> List[Optional[str]]:
+        return [item.id for item in self.data]
+
+    @property
+    def output(self) -> List[Dict[str, Any]]:
+        return [item.output for item in self.data]
+
+    def get_batch_data(self, attr_name: str, batch_size: int) -> Generator[List[Any], None, None]:
+        """Get an attribute of dataset items in batch."""
+        for i in range(0, len(self.data), batch_size):
+            batch_items = self.data[i : i + batch_size]
+            yield [item[attr_name] for item in batch_items]
+
+    def __getattr__(self, attr_name: str) -> List[Any]:
+        return [item.__getattr__(attr_name) for item in self.data]
+
+    def get_attr_data(self, attr_name: str) -> List[Any]:
+        """For the attributes constructed later (not implemented using property),
+        obtain a list of this attribute in the entire dataset.
+        """
+        return [item[attr_name] for item in self.data]
+
+    def __getitem__(self, index: int) -> Item:
+        return self.data[index]
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def save(self, save_path: str) -> None:
+        """Save the dataset into the original format."""
+
+        save_data = [item.to_dict() for item in self.data]
+
+        with open(save_path, "w", encoding="utf-8") as f:
+            json.dump(save_data, f, indent=4)
+
+    def __str__(self) -> str:
+        """Return a string representation of the dataset with a summary of items."""
+        return f"Dataset '{self.dataset_name}' with {len(self)} items"

{flashrag-dev-0.1.3.dev20241108 → flashrag-dev-0.1.3.dev20241111}/flashrag/dataset/utils.py

@@ -1,6 +1,22 @@
+from typing import Dict, Any, Union
+import numpy as np
 from flashrag.dataset import Dataset
 
 
+def convert_numpy(obj: Union[Dict, list, np.ndarray, np.generic]) -> Any:
+    """Recursively convert numpy objects in nested dictionaries or lists to native Python types."""
+    if isinstance(obj, dict):
+        return {k: convert_numpy(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [convert_numpy(i) for i in obj]
+    elif isinstance(obj, np.ndarray):
+        return obj.tolist()  # Convert numpy arrays to lists
+    elif isinstance(obj, (np.integer, np.floating)):
+        return obj.item()  # Convert numpy scalars to native Python scalars
+    else:
+        return obj  # Return the object as-is if it's neither a dict, list, nor numpy type
+
+
 def filter_dataset(dataset: Dataset, filter_func=None):
     if filter_func is None:
         return dataset
@@ -17,28 +33,20 @@ def split_dataset(dataset: Dataset, split_symbol: list):
     data = dataset.data
     data_split = {symbol: [] for symbol in set(split_symbol)}
     for symbol in set(split_symbol):
-        symbol_data = [
-            x for x, x_symbol in zip(data, split_symbol) if x_symbol == symbol
-        ]
+        symbol_data = [x for x, x_symbol in zip(data, split_symbol) if x_symbol == symbol]
         data_split[symbol] = Dataset(config=dataset.config, data=symbol_data)
 
     return data_split
 
 
 def merge_dataset(dataset_split: dict, split_symbol: list):
-    assert len(split_symbol) == sum(
-        [len(data) for data in dataset_split.values()]
-    )
-    dataset_split_iter = {
-        symbol: iter(dataset.data) for symbol, dataset in dataset_split.items()
-    }
+    assert len(split_symbol) == sum([len(data) for data in dataset_split.values()])
+    dataset_split_iter = {symbol: iter(dataset.data) for symbol, dataset in dataset_split.items()}
 
     final_data = []
     for item_symbol in split_symbol:
         final_data.append(next(dataset_split_iter[item_symbol]))
-    final_dataset = Dataset(
-        config=list(dataset_split.values())[0].config, data=final_data
-    )
+    final_dataset = Dataset(config=list(dataset_split.values())[0].config, data=final_data)
 
     return final_dataset
 

flashrag-dev-0.1.3.dev20241111/flashrag/generator/__init__.py

@@ -0,0 +1,3 @@
+from flashrag.generator.generator import *
+from flashrag.generator.openai_generator import *
+from flashrag.generator.utils import *

{flashrag-dev-0.1.3.dev20241108 → flashrag-dev-0.1.3.dev20241111}/flashrag/generator/generator.py

@@ -1,5 +1,6 @@
 from typing import List
 from copy import deepcopy
+import warnings
 from tqdm import tqdm
 from tqdm.auto import trange
 import numpy as np
@@ -11,6 +12,7 @@ from transformers import (
     BartForConditionalGeneration,
     AutoConfig,
 )
+from flashrag.generator.utils import resolve_max_tokens
 
 
 class BaseGenerator:
@@ -53,15 +55,11 @@ class EncoderDecoderGenerator(BaseGenerator):
 
                 self.model = FiDT5.from_pretrained(self.model_path)
             else:
-                self.model = T5ForConditionalGeneration.from_pretrained(
-                    self.model_path
-                )
+                self.model = T5ForConditionalGeneration.from_pretrained(self.model_path)
         else:
             if self.fid:
                 assert False, "FiD only support T5"
-            self.model = BartForConditionalGeneration.from_pretrained(
-                self.model_path
-            )
+            self.model = BartForConditionalGeneration.from_pretrained(self.model_path)
         self.model.cuda()
         self.model.eval()
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
@@ -108,27 +106,14 @@ class EncoderDecoderGenerator(BaseGenerator):
             ]
             generation_params["stopping_criteria"] = stopping_criteria
 
-        max_tokens = params.pop("max_tokens", None) or params.pop(
-            "max_new_tokens", None
-        )
-        if max_tokens is not None:
-            generation_params["max_new_tokens"] = max_tokens
-        else:
-            generation_params["max_new_tokens"] = generation_params.get(
-                "max_new_tokens", generation_params.pop("max_tokens", None)
-            )
-        generation_params.pop("max_tokens", None)
+        generation_params = resolve_max_tokens(params, generation_params, prioritize_new_tokens=True)
 
         responses = []
-        for idx in trange(
-            0, len(input_list), batch_size, desc="Generation process: "
-        ):
+        for idx in trange(0, len(input_list), batch_size, desc="Generation process: "):
             batched_prompts = input_list[idx : idx + batch_size]
             if self.fid:
                 # assume each input in input_list is a list, contains K string
-                input_ids, attention_mask = self.encode_passages(
-                    batched_prompts
-                )
+                input_ids, attention_mask = self.encode_passages(batched_prompts)
                 inputs = {
                     "input_ids": input_ids.to(self.device),
                     "attention_mask": attention_mask.to(self.device),
@@ -173,11 +158,7 @@ class VLLMGenerator(BaseGenerator):
         else:
             tensor_parallel_size = self.gpu_num
 
-        self.lora_path = (
-            None
-            if "generator_lora_path" not in config
-            else config["generator_lora_path"]
-        )
+        self.lora_path = None if "generator_lora_path" not in config else config["generator_lora_path"]
         self.use_lora = False
         if self.lora_path is not None:
             self.use_lora = True
@@ -197,9 +178,7 @@ class VLLMGenerator(BaseGenerator):
                 gpu_memory_utilization=gpu_memory_utilization,
                 max_logprobs=32016,
             )
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            self.model_path, trust_remote_code=True
-        )
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True)
 
     @torch.inference_mode(mode=True)
     def generate(
@@ -219,23 +198,16 @@ class VLLMGenerator(BaseGenerator):
         if "do_sample" in generation_params:
             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
-        )
-        if max_tokens is not None:
-            generation_params["max_tokens"] = max_tokens
-        else:
-            generation_params["max_tokens"] = generation_params.get(
-                "max_tokens", generation_params.pop("max_new_tokens", None)
-            )
-        generation_params.pop("max_new_tokens", None)
+                generation_params["temperature"] = 0
+        generation_params["seed"] = self.config["seed"]
+
+        # handle param conflict
+        generation_params = resolve_max_tokens(params, generation_params, prioritize_new_tokens=False)
 
         # fix for llama3
         if "stop" in generation_params:
             generation_params["stop"].append("<|eot_id|>")
-            generation_params['include_stop_str_in_output'] = True
+            generation_params["include_stop_str_in_output"] = True
         else:
             generation_params["stop"] = ["<|eot_id|>"]
 
@@ -265,12 +237,7 @@ class VLLMGenerator(BaseGenerator):
             scores = []
             for output in outputs:
                 logprobs = output.outputs[0].logprobs
-                scores.append(
-                    [
-                        np.exp(list(score_dict.values())[0].logprob)
-                        for score_dict in logprobs
-                    ]
-                )
+                scores.append([np.exp(list(score_dict.values())[0].logprob) for score_dict in logprobs])
             return base_output, scores
         else:
             return base_output
@@ -282,11 +249,7 @@ class HFCausalLMGenerator(BaseGenerator):
     def __init__(self, config, model=None):
         super().__init__(config)
         self.config = config
-        lora_path = (
-            None
-            if "generator_lora_path" not in config
-            else config["generator_lora_path"]
-        )
+        lora_path = None if "generator_lora_path" not in config else config["generator_lora_path"]
         self.model, self.tokenizer = self._load_model(model=model)
         self.use_lora = False
         if lora_path is not None:
@@ -305,28 +268,24 @@ class HFCausalLMGenerator(BaseGenerator):
         else:
             model.cuda()
         model.eval()
-        tokenizer = AutoTokenizer.from_pretrained(
-            self.model_path, trust_remote_code=True
-        )
+        tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True)
         if "qwen" not in self.model_name:
             tokenizer.pad_token = tokenizer.eos_token
         tokenizer.padding_side = "left"
 
         return model, tokenizer
 
-    def add_new_tokens(
-        self, token_embedding_path, token_name_func=lambda idx: f"[ref{idx+1}]"
-    ):
+    def add_new_tokens(self, token_embedding_path, token_name_func=lambda idx: f"[ref{idx+1}]"):
         del self.model
         self.model = AutoModelForCausalLM.from_pretrained(
-                self.model_path,
-                trust_remote_code=True,
+            self.model_path,
+            trust_remote_code=True,
         )
         # get original embedding weight matrix
         embedding_layer = self.model.get_input_embeddings()
         embedding_weights = embedding_layer.weight
         original_vocab_size, embedding_dim = embedding_weights.shape
-    
+
         new_tokens_weights = torch.load(token_embedding_path)
         new_tokens_length = new_tokens_weights.shape[0]
 
@@ -386,16 +345,7 @@ class HFCausalLMGenerator(BaseGenerator):
             ]
             generation_params["stopping_criteria"] = stopping_criteria
 
-        max_tokens = params.pop("max_tokens", None) or params.pop(
-            "max_new_tokens", None
-        )
-        if max_tokens is not None:
-            generation_params["max_new_tokens"] = max_tokens
-        else:
-            generation_params["max_new_tokens"] = generation_params.get(
-                "max_new_tokens", generation_params.pop("max_tokens", None)
-            )
-        generation_params.pop("max_tokens", None)
+        generation_params = resolve_max_tokens(params, generation_params, prioritize_new_tokens=True)
 
         # set eos token for llama
         if "llama" in self.model_name.lower():
@@ -413,9 +363,7 @@ class HFCausalLMGenerator(BaseGenerator):
         generated_token_ids = []
         generated_token_logits = []
 
-        for idx in trange(
-            0, len(input_list), batch_size, desc="Generation process: "
-        ):
+        for idx in trange(0, len(input_list), batch_size, desc="Generation process: "):
             torch.cuda.empty_cache()
             batched_prompts = input_list[idx : idx + batch_size]
             inputs = self.tokenizer(
@@ -435,12 +383,7 @@ class HFCausalLMGenerator(BaseGenerator):
             generated_ids = outputs.sequences
             logits = torch.stack(outputs.scores, dim=1).softmax(-1)
             generated_ids = generated_ids[:, inputs["input_ids"].shape[-1] :]
-            gen_score = (
-                torch.gather(logits, 2, generated_ids[:, :, None])
-                .squeeze(-1)
-                .cpu()
-                .tolist()
-            )
+            gen_score = torch.gather(logits, 2, generated_ids[:, :, None]).squeeze(-1).cpu().tolist()
             scores.extend(gen_score)
 
             # get additinoal info
@@ -448,26 +391,15 @@ class HFCausalLMGenerator(BaseGenerator):
                 batch_generated_token_ids = generated_ids.detach().cpu()
                 batch_generated_token_logits = (
                     torch.cat(
-                        [
-                            token_scores.unsqueeze(1)
-                            for token_scores in outputs.scores
-                        ],
+                        [token_scores.unsqueeze(1) for token_scores in outputs.scores],
                         dim=1,
                     )
                     .detach()
                     .cpu()
                 )
-                if (
-                    batch_generated_token_ids.shape[1]
-                    < generation_params["max_new_tokens"]
-                ):
-                    real_batch_size, num_generated_tokens = (
-                        batch_generated_token_ids.shape
-                    )
-                    padding_length = (
-                        generation_params["max_new_tokens"]
-                        - num_generated_tokens
-                    )
+                if batch_generated_token_ids.shape[1] < generation_params["max_new_tokens"]:
+                    real_batch_size, num_generated_tokens = batch_generated_token_ids.shape
+                    padding_length = generation_params["max_new_tokens"] - num_generated_tokens
                     padding_token_ids = torch.zeros(
                         (real_batch_size, padding_length),
                         dtype=batch_generated_token_ids.dtype,
@@ -480,9 +412,7 @@ class HFCausalLMGenerator(BaseGenerator):
                         ),
                         dtype=batch_generated_token_logits.dtype,
                     )
-                    batch_generated_token_ids = torch.cat(
-                        [batch_generated_token_ids, padding_token_ids], dim=1
-                    )
+                    batch_generated_token_ids = torch.cat([batch_generated_token_ids, padding_token_ids], dim=1)
                     batch_generated_token_logits = torch.cat(
                         [batch_generated_token_logits, padding_token_logits],
                         dim=1,
@@ -569,11 +499,7 @@ class FastChatGenerator(HFCausalLMGenerator):
         def get_gpu_memory(max_gpus=None):
             """Get available memory for each GPU."""
             gpu_memory = []
-            num_gpus = (
-                torch.cuda.device_count()
-                if max_gpus is None
-                else min(max_gpus, torch.cuda.device_count())
-            )
+            num_gpus = torch.cuda.device_count() if max_gpus is None else min(max_gpus, torch.cuda.device_count())
             for gpu_id in range(num_gpus):
                 with torch.cuda.device(gpu_id):
                     device = torch.cuda.current_device()
@@ -594,10 +520,7 @@ class FastChatGenerator(HFCausalLMGenerator):
             max_gpu_memory = None
             if self.gpu_num != 1:
                 available_gpu_memory = get_gpu_memory(self.gpu_num)
-                max_gpu_memory = (
-                    str(int(min(available_gpu_memory) * gpu_memory_utilization))
-                    + "GiB"
-                )
+                max_gpu_memory = str(int(min(available_gpu_memory) * gpu_memory_utilization)) + "GiB"
 
             model, tokenizer = load_model(
                 self.model_path,
@@ -611,14 +534,10 @@ class FastChatGenerator(HFCausalLMGenerator):
 
         else:
             model.cuda()
-            tokenizer = AutoTokenizer.from_pretrained(
-                self.model_path, trust_remote_code=True
-            )
+            tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True)
         model.eval()
 
-        tokenizer = AutoTokenizer.from_pretrained(
-            self.model_path, trust_remote_code=True
-        )
+        tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True)
         if "qwen" not in self.model_name:
             tokenizer.pad_token = tokenizer.eos_token
         tokenizer.padding_side = "left"

flashrag-dev-0.1.3.dev20241111/flashrag/generator/utils.py

@@ -0,0 +1,67 @@
+import warnings
+
+
+def resolve_max_tokens(params: dict, generation_params: dict, prioritize_new_tokens: bool = False) -> dict:
+    """
+    Resolve and validate max_tokens parameters from both params and generation_params.
+
+    Args:
+        params: Dictionary containing user-provided parameters
+        generation_params: Dictionary containing generation-specific parameters
+        prioritize_new_tokens: If True, max_new_tokens takes precedence over max_tokens
+                             If False, max_tokens takes precedence (default behavior)
+
+    Returns:
+        Updated generation_params dictionary
+    """
+
+    def get_token_params(param_dict: dict) -> tuple:
+        """Extract max_tokens and max_new_tokens from a parameter dictionary."""
+        return (param_dict.pop("max_tokens", None), param_dict.pop("max_new_tokens", None))
+
+    def resolve_tokens(max_tokens: int, max_new_tokens: int) -> int:
+        """
+        Resolve between max_tokens and max_new_tokens values based on priority.
+        Returns the resolved token value or None if no valid value found.
+        """
+        # If either value is None, return the non-None value
+        if max_tokens is None:
+            return max_new_tokens
+        if max_new_tokens is None:
+            return max_tokens
+
+        # Both values exist but are different
+        if max_tokens != max_new_tokens:
+            if prioritize_new_tokens:
+                warnings.warn(
+                    f"max_tokens ({max_tokens}) and max_new_tokens ({max_new_tokens}) "
+                    f"are different. Using max_new_tokens value as it has priority."
+                )
+                return max_new_tokens
+            else:
+                warnings.warn(
+                    f"max_tokens ({max_tokens}) and max_new_tokens ({max_new_tokens}) "
+                    f"are different. Using max_tokens value as it has priority."
+                )
+                return max_tokens
+
+        # Both values are equal
+        return max_tokens
+
+    # Try to resolve from params first, then fall back to generation_params
+    max_tokens, max_new_tokens = get_token_params(params)
+    final_max_tokens = resolve_tokens(max_tokens, max_new_tokens)
+
+    # If no valid tokens found in params, try generation_params
+    if final_max_tokens is None:
+        max_tokens, max_new_tokens = get_token_params(generation_params)
+        final_max_tokens = resolve_tokens(max_tokens, max_new_tokens)
+
+    generation_params.pop("max_new_tokens", None)
+    generation_params.pop("max_tokens", None)
+    if final_max_tokens is not None:
+        if prioritize_new_tokens:
+            generation_params["max_new_tokens"] = final_max_tokens
+        else:
+            generation_params["max_tokens"] = final_max_tokens
+    return generation_params

{flashrag-dev-0.1.3.dev20241108 → flashrag-dev-0.1.3.dev20241111}/flashrag/retriever/encoder.py

@@ -5,6 +5,22 @@ from flashrag.retriever.utils import load_model, pooling, parse_query
 
 
 class Encoder:
+    """
+    Encoder class for encoding queries using a specified model.
+
+    Attributes:
+        model_name (str): The name of the model.
+        model_path (str): The path to the model.
+        pooling_method (str): The method used for pooling.
+        max_length (int): The maximum length of the input sequences.
+        use_fp16 (bool): Whether to use FP16 precision.
+        instruction (str): Additional instructions for parsing queries.
+
+    Methods:
+        encode(query_list: List[str], is_query=True) -> np.ndarray:
+            Encodes a list of queries into embeddings.
+    """
+
     def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16, instruction):
         self.model_name = model_name
         self.model_path = model_path
@@ -15,7 +31,7 @@ class Encoder:
 
         self.model, self.tokenizer = load_model(model_path=model_path, use_fp16=use_fp16)
 
-    @torch.inference_mode(mode=True)
+    @torch.inference_mode()
     def encode(self, query_list: List[str], is_query=True) -> np.ndarray:
         query_list = parse_query(self.model_name, query_list, self.instruction)
 
@@ -44,6 +60,23 @@ class Encoder:
 
 
 class STEncoder:
+    """
+    STEncoder class for encoding queries using SentenceTransformers.
+
+    Attributes:
+        model_name (str): The name of the model.
+        model_path (str): The path to the model.
+        max_length (int): The maximum length of the input sequences.
+        use_fp16 (bool): Whether to use FP16 precision.
+        instruction (str): Additional instructions for parsing queries.
+
+    Methods:
+        encode(query_list: List[str], batch_size=64, is_query=True) -> np.ndarray:
+            Encodes a list of queries into embeddings.
+        multi_gpu_encode(query_list: List[str], is_query=True, batch_size=None) -> np.ndarray:
+            Encodes a list of queries into embeddings using multiple GPUs.
+    """
+
     def __init__(self, model_name, model_path, max_length, use_fp16, instruction):
         import torch
         from sentence_transformers import SentenceTransformer
@@ -54,10 +87,10 @@ class STEncoder:
         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}
+            model_path, trust_remote_code=True, model_kwargs={"torch_dtype": torch.float16 if use_fp16 else torch.float}
         )
 
-    @torch.inference_mode(mode=True)
+    @torch.inference_mode()
     def encode(self, query_list: List[str], batch_size=64, is_query=True) -> np.ndarray:
         query_list = parse_query(self.model_name, query_list, self.instruction)
         query_emb = self.model.encode(
@@ -67,14 +100,19 @@ class STEncoder:
 
         return query_emb
 
-    @torch.inference_mode(mode=True)
+    @torch.inference_mode()
     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, 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,show_progress_bar=True
+            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")
+        query_emb = query_emb.astype(np.float32, order="C")
 
         return query_emb

flashrag-dev-0.1.3.dev20241111/flashrag/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.3dev20241111"

{flashrag-dev-0.1.3.dev20241108 → flashrag-dev-0.1.3.dev20241111/flashrag_dev.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag-dev
-Version: 0.1.3.dev20241108
+Version: 0.1.3.dev20241111
 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

{flashrag-dev-0.1.3.dev20241108 → flashrag-dev-0.1.3.dev20241111}/flashrag_dev.egg-info/SOURCES.txt

@@ -20,6 +20,7 @@ flashrag/generator/fid.py
 flashrag/generator/generator.py
 flashrag/generator/openai_generator.py
 flashrag/generator/stop_word_criteria.py
+flashrag/generator/utils.py
 flashrag/judger/__init__.py
 flashrag/judger/judger.py
 flashrag/pipeline/__init__.py

flashrag-dev-0.1.3.dev20241108/flashrag/dataset/dataset.py

@@ -1,160 +0,0 @@
-import os
-import json
-import random
-import numpy as np
-
-
-class Item:
-    r"""A container class used to store and manipulate a sample within a dataset.
-    Information related to this sample during training/inference will be stored in ```self.output```.
-    Each attribute of this class can be used like a dict key(also for key in ```self.output```).
-
-    """
-
-    def __init__(self, item_dict):
-        self.id = item_dict.get("id", None)
-        self.question = item_dict.get("question", None)
-        self.golden_answers = item_dict.get("golden_answers", [])
-        self.choices = item_dict.get("choices", [])
-        self.metadata = item_dict.get("metadata", {})
-        self.output = item_dict.get("output", {})
-
-    def update_output(self, key, value):
-        r"""Update the output dict and keep a key in self.output can be used as an attribute."""
-        if key in ["id", "question", "golden_answers", "output"]:
-            raise AttributeError(f"{key} should not be changed")
-        else:
-            self.output[key] = value
-
-    def update_evaluation_score(self, metric_name, metric_score):
-        r"""Update the evaluation score of this sample for a metric."""
-        if "metric_score" not in self.output:
-            self.output["metric_score"] = {}
-        self.output["metric_score"][metric_name] = metric_score
-
-    def __getattr__(self, attr_name):
-        if attr_name in ["id", "question", "golden_answers", "metadata", "output", "choices"]:
-            return super().__getattribute__(attr_name)
-        else:
-            output = super().__getattribute__("output")
-            if attr_name in output:
-                return output[attr_name]
-            else:
-                raise AttributeError(f"Attribute `{attr_name}` not found")
-
-    def to_dict(self):
-        r"""Convert all information within the data sample into a dict. Information generated
-        during the inference will be saved into output field.
-
-        """
-        for k, v in self.output.items():
-            if isinstance(k, np.ndarray):
-                self.output[k] = v.tolist()
-        output = {
-            "id": self.id,
-            "question": self.question,
-            "golden_answers": self.golden_answers,
-            "output": self.output,
-        }
-        if self.metadata != {}:
-            output["metadata"] = self.metadata
-
-        return output
-
-
-class Dataset:
-    """A container class used to store the whole dataset. Inside the class, each data sample will be stored
-    in ```Item``` class.
-    The properties of the dataset represent the list of attributes corresponding to each item in the dataset.
-    """
-
-    def __init__(self, config=None, dataset_path=None, data=None, sample_num=None, random_sample=False):
-        self.config = config
-        self.dataset_name = config["dataset_name"]
-        self.dataset_path = dataset_path
-
-        self.sample_num = sample_num
-        self.random_sample = random_sample
-
-        if data is None:
-            self.data = self._load_data(self.dataset_name, self.dataset_path)
-        else:
-            self.data = data
-
-    def _load_data(self, dataset_name, dataset_path):
-        """Load data from the provided dataset_path or directly download the file(TODO)."""
-
-        if not os.path.exists(dataset_path):
-            # TODO: auto download: self._download(dataset_name, dataset_path)
-            pass
-
-        data = []
-        with open(dataset_path, "r", encoding="utf-8") as f:
-            for line in f:
-                item_dict = json.loads(line)
-                item = Item(item_dict)
-                data.append(item)
-        if self.sample_num is not None:
-            if self.random_sample:
-                print(f"Random sample {self.sample_num} items in test set.")
-                data = random.sample(data, self.sample_num)
-            else:
-                data = data[: self.sample_num]
-
-        return data
-
-    def update_output(self, key, value_list):
-        """Update the overall output field for each sample in the dataset."""
-
-        assert len(self.data) == len(value_list)
-        for item, value in zip(self.data, value_list):
-            item.update_output(key, value)
-
-    @property
-    def question(self):
-        return [item.question for item in self.data]
-
-    @property
-    def golden_answers(self):
-        return [item.golden_answers for item in self.data]
-
-    @property
-    def id(self):
-        return [item.id for item in self.data]
-
-    @property
-    def output(self):
-        return [item.output for item in self.data]
-
-    def get_batch_data(self, attr_name: str, batch_size: int):
-        """Get an attribute of dataset items in batch."""
-
-        for i in range(0, len(self.data), batch_size):
-            batch_items = self.data[i : i + batch_size]
-            yield [item[attr_name] for item in batch_items]
-
-    def __getattr__(self, attr_name):
-        return [item.__getattr__(attr_name) for item in self.data]
-
-    def get_attr_data(self, attr_name):
-        """For the attributes constructed later (not implemented using property),
-        obtain a list of this attribute in the entire dataset.
-        """
-        return [item[attr_name] for item in self.data]
-
-    def __getitem__(self, index):
-        return self.data[index]
-
-    def __len__(self):
-        return len(self.data)
-
-    def save(self, save_path):
-        """Save the dataset into the original format."""
-
-        def convert_to_float(d):
-            return {k: (v.item() if isinstance(v, np.generic) else v) for k, v in d.items()}
-
-        save_data = [convert_to_float(item.to_dict()) for item in self.data]
-
-        with open(save_path, "w") as f:
-            json.dump(save_data, f, indent=4)

flashrag-dev-0.1.3.dev20241108/flashrag/generator/__init__.py

@@ -1,2 +0,0 @@
-from flashrag.generator.generator import *
-from flashrag.generator.openai_generator import *

flashrag-dev-0.1.3.dev20241108/flashrag/version.py

@@ -1 +0,0 @@
-__version__ = "0.1.3dev20241108"