flashrag-dev 0.1.1__py3-none-any.whl

flashrag/config/__init__.py

@@ -0,0 +1,2 @@
+from flashrag.config.config import Config
+

flashrag/config/config.py

@@ -0,0 +1,219 @@
+import re
+import os
+import yaml
+import random
+import datetime
+
+
+class Config:
+    def __init__(self, config_file_path=None, config_dict={}):
+
+        self.yaml_loader = self._build_yaml_loader()
+        self.file_config = self._load_file_config(config_file_path)
+        self.variable_config = config_dict
+
+        self.external_config = self._merge_external_config()
+
+        self.internal_config = self._get_internal_config()
+
+        self.final_config = self._get_final_config()
+
+        self._check_final_config()
+        self._set_additional_key()
+
+        self._init_device()
+        self._set_seed()
+        self._prepare_dir()
+
+    def _build_yaml_loader(self):
+        loader = yaml.FullLoader
+        loader.add_implicit_resolver(
+            "tag:yaml.org,2002:float",
+            re.compile(
+                """^(?:
+             [-+]?(?:[0-9][0-9_]*)\\.[0-9_]*(?:[eE][-+]?[0-9]+)?
+            |[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+)
+            |\\.[0-9_]+(?:[eE][-+][0-9]+)?
+            |[-+]?[0-9][0-9_]*(?::[0-5]?[0-9])+\\.[0-9_]*
+            |[-+]?\\.(?:inf|Inf|INF)
+            |\\.(?:nan|NaN|NAN))$""",
+                re.X,
+            ),
+            list("-+0123456789."),
+        )
+        return loader
+
+    def _load_file_config(self, config_file_path: str):
+        file_config = dict()
+        if config_file_path:
+            with open(config_file_path, "r", encoding="utf-8") as f:
+                file_config.update(yaml.load(f.read(), Loader=self.yaml_loader))
+        return file_config
+
+    @staticmethod
+    def _update_dict(old_dict: dict, new_dict: dict):
+        # Update the original update method of the dictionary:
+        # If there is the same key in `old_dict` and `new_dict`, and value is of type dict, update the key in dict
+
+        same_keys = []
+        for key, value in new_dict.items():
+            if key in old_dict and isinstance(value, dict):
+                same_keys.append(key)
+        for key in same_keys:
+            old_item = old_dict[key]
+            new_item = new_dict[key]
+            old_item.update(new_item)
+            new_dict[key] = old_item
+
+        old_dict.update(new_dict)
+        return old_dict
+
+    def _merge_external_config(self):
+        external_config = dict()
+        external_config = self._update_dict(external_config, self.file_config)
+        external_config = self._update_dict(external_config, self.variable_config)
+
+        return external_config
+
+    def _get_internal_config(self):
+        current_path = os.path.dirname(os.path.realpath(__file__))
+        init_config_path = os.path.join(current_path, "basic_config.yaml")
+        internal_config = self._load_file_config(init_config_path)
+
+        return internal_config
+
+    def _get_final_config(self):
+        final_config = dict()
+        final_config = self._update_dict(final_config, self.internal_config)
+        final_config = self._update_dict(final_config, self.external_config)
+
+        return final_config
+
+    def _check_final_config(self):
+        # check split
+        split = self.final_config["split"]
+        if split is None:
+            split = ["train", "dev", "test"]
+        if isinstance(split, str):
+            split = [split]
+        self.final_config["split"] = split
+
+    def _init_device(self):
+        gpu_id = self.final_config["gpu_id"]
+        if gpu_id is not None:
+            os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
+            import torch
+
+            self.final_config["device"] = torch.device("cuda")
+        else:
+            import torch
+
+            self.final_config["device"] = torch.device("cpu")
+
+    def _set_additional_key(self):
+        # set dataset
+        dataset_name = self.final_config["dataset_name"]
+        data_dir = self.final_config["data_dir"]
+        self.final_config["dataset_path"] = os.path.join(data_dir, dataset_name)
+
+        # set model path
+        retrieval_method = self.final_config["retrieval_method"]
+        model2path = self.final_config["model2path"]
+        model2pooling = self.final_config["model2pooling"]
+        method2index = self.final_config["method2index"]
+
+        generator_model = self.final_config["generator_model"]
+
+        if self.final_config["index_path"] is None:
+            try:
+                self.final_config["index_path"] = method2index[retrieval_method]
+            except:
+                print("Index is empty!!")
+                assert False
+
+        self.final_config["retrieval_model_path"] = model2path.get(retrieval_method, retrieval_method)
+        # TODO: not support when `retrieval_model` is path
+
+        def set_pooling_method(method, model2pooling):
+            for key, value in model2pooling.items():
+                if key.lower() in method.lower():
+                    return value
+            return "mean"
+
+        if self.final_config.get("retrieval_pooling_method") is None:
+            self.final_config["retrieval_pooling_method"] = set_pooling_method(retrieval_method, model2pooling)
+
+        rerank_model_name = self.final_config["rerank_model_name"]
+        if self.final_config.get("rerank_model_path") is None:
+            if rerank_model_name is not None:
+                self.final_config["rerank_model_path"] = model2path.get(rerank_model_name, rerank_model_name)
+        if self.final_config["rerank_pooling_method"] is None:
+            if rerank_model_name is not None:
+                self.final_config["rerank_pooling_method"] = set_pooling_method(rerank_model_name, model2pooling)
+
+        if self.final_config.get("generator_model_path") is None:
+            self.final_config["generator_model_path"] = model2path.get(generator_model, generator_model)
+
+        if "refiner_name" in self.final_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)
+
+        # set model path in metric setting
+        metric_setting = self.final_config["metric_setting"]
+        metric_tokenizer_name = metric_setting.get("tokenizer_name", None)
+        from flashrag.utils.constants import OPENAI_MODEL_DICT
+
+        if metric_tokenizer_name not in OPENAI_MODEL_DICT:
+            metric_tokenizer_name = model2path.get(metric_tokenizer_name, metric_tokenizer_name)
+            metric_setting["tokenizer_name"] = metric_tokenizer_name
+            self.final_config["metric_setting"] = metric_setting
+
+    def _prepare_dir(self):
+        save_note = self.final_config["save_note"]
+        current_time = datetime.datetime.now()
+        self.final_config["save_dir"] = os.path.join(
+            self.final_config["save_dir"],
+            f"{self.final_config['dataset_name']}_{current_time.strftime('%Y_%m_%d_%H_%M')}_{save_note}",
+        )
+        os.makedirs(self.final_config["save_dir"], exist_ok=True)
+        # save config parameters
+        config_save_path = os.path.join(self.final_config["save_dir"], "config.yaml")
+        with open(config_save_path, "w") as f:
+            yaml.dump(self.final_config, f)
+
+    def _set_seed(self):
+        import torch
+        import numpy as np
+
+        seed = self.final_config["seed"]
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.deterministic = True
+
+    def __setitem__(self, key, value):
+        if not isinstance(key, str):
+            raise TypeError("index must be a str.")
+        self.final_config[key] = value
+
+    def __getattr__(self, item):
+        if "final_config" not in self.__dict__:
+            raise AttributeError("'Config' object has no attribute 'final_config'")
+        if item in self.final_config:
+            return self.final_config[item]
+        raise AttributeError(f"'Config' object has no attribute '{item}'")
+
+    def __getitem__(self, item):
+        return self.final_config.get(item)
+
+    def __contains__(self, key):
+        if not isinstance(key, str):
+            raise TypeError("index must be a str.")
+        return key in self.final_config
+
+    def __repr__(self):
+        return self.final_config.__str__()

flashrag/dataset/__init__.py

@@ -0,0 +1,2 @@
+from flashrag.dataset.dataset import *
+from flashrag.dataset.utils import *

flashrag/dataset/dataset.py

@@ -0,0 +1,160 @@
+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/dataset/utils.py

@@ -0,0 +1,60 @@
+from flashrag.dataset import Dataset
+
+
+def filter_dataset(dataset: Dataset, filter_func=None):
+    if filter_func is None:
+        return dataset
+    data = dataset.data
+    for item in data:
+        if not filter_func(item):
+            data.remove(item)
+    return Dataset(config=dataset.config, data=data)
+
+
+def split_dataset(dataset: Dataset, split_symbol: list):
+    assert len(split_symbol) == len(dataset)
+
+    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
+        ]
+        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()
+    }
+
+    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
+    )
+
+    return final_dataset
+
+
+def get_batch_dataset(dataset: Dataset, batch_size=16):
+    data = dataset.data
+    for idx in range(0, len(data), batch_size):
+        batched_data = data[idx : idx + batch_size]
+        batch_dataset = Dataset(config=dataset.config, data=batched_data)
+        yield batch_dataset
+
+
+def merge_batch_dataset(dataset_list: Dataset):
+    dataset = dataset_list[0]
+    total_data = []
+    for batch_dataset in dataset_list:
+        total_data.extend(batch_dataset.data)
+    dataset = Dataset(config=dataset.config, data=total_data)
+    return dataset

flashrag/evaluator/__init__.py

@@ -0,0 +1,2 @@
+from flashrag.evaluator.evaluator import *
+from flashrag.evaluator.metrics import *

flashrag/evaluator/_bleu.py

@@ -0,0 +1,209 @@
+# Source: https://github.com/mjpost/sacrebleu/blob/master/sacrebleu/tokenizers/tokenizer_13a.py
+# Copyright 2020 SacreBLEU Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+from functools import lru_cache
+
+
+class BaseTokenizer:
+    """A base dummy tokenizer to derive from."""
+
+    def signature(self):
+        """
+        Returns a signature for the tokenizer.
+        :return: signature string
+        """
+        return "none"
+
+    def __call__(self, line):
+        """
+        Tokenizes an input line with the tokenizer.
+        :param line: a segment to tokenize
+        :return: the tokenized line
+        """
+        return line
+
+
+class TokenizerRegexp(BaseTokenizer):
+    def signature(self):
+        return "re"
+
+    def __init__(self):
+        self._re = [
+            # language-dependent part (assuming Western languages)
+            (re.compile(r"([\{-\~\[-\` -\&\(-\+\:-\@\/])"), r" \1 "),
+            # tokenize period and comma unless preceded by a digit
+            (re.compile(r"([^0-9])([\.,])"), r"\1 \2 "),
+            # tokenize period and comma unless followed by a digit
+            (re.compile(r"([\.,])([^0-9])"), r" \1 \2"),
+            # tokenize dash when preceded by a digit
+            (re.compile(r"([0-9])(-)"), r"\1 \2 "),
+            # one space only between words
+            # NOTE: Doing this in Python (below) is faster
+            # (re.compile(r'\s+'), r' '),
+        ]
+
+    @lru_cache(maxsize=2**16)
+    def __call__(self, line):
+        """Common post-processing tokenizer for `13a` and `zh` tokenizers.
+        :param line: a segment to tokenize
+        :return: the tokenized line
+        """
+        for _re, repl in self._re:
+            line = _re.sub(repl, line)
+
+        # no leading or trailing spaces, single space within words
+        # return ' '.join(line.split())
+        # This line is changed with regards to the original tokenizer (seen above) to return individual words
+        return line.split()
+
+
+class Tokenizer13a(BaseTokenizer):
+    def signature(self):
+        return "13a"
+
+    def __init__(self):
+        self._post_tokenizer = TokenizerRegexp()
+
+    @lru_cache(maxsize=2**16)
+    def __call__(self, line):
+        """Tokenizes an input line using a relatively minimal tokenization
+        that is however equivalent to mteval-v13a, used by WMT.
+        :param line: a segment to tokenize
+        :return: the tokenized line
+        """
+
+        # language-independent part:
+        line = line.replace("<skipped>", "")
+        line = line.replace("-\n", "")
+        line = line.replace("\n", " ")
+
+        if "&" in line:
+            line = line.replace("&quot;", '"')
+            line = line.replace("&amp;", "&")
+            line = line.replace("&lt;", "<")
+            line = line.replace("&gt;", ">")
+
+        return self._post_tokenizer(f" {line} ")
+
+
+# Copyright 2017 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Python implementation of BLEU and smooth-BLEU.
+
+This module provides a Python implementation of BLEU and smooth-BLEU.
+Smooth BLEU is computed following the method outlined in the paper:
+Chin-Yew Lin, Franz Josef Och. ORANGE: a method for evaluating automatic
+evaluation metrics for machine translation. COLING 2004.
+"""
+
+import collections
+import math
+
+
+def _get_ngrams(segment, max_order):
+    """Extracts all n-grams upto a given maximum order from an input segment.
+
+    Args:
+      segment: text segment from which n-grams will be extracted.
+      max_order: maximum length in tokens of the n-grams returned by this
+          methods.
+
+    Returns:
+      The Counter containing all n-grams upto max_order in segment
+      with a count of how many times each n-gram occurred.
+    """
+    ngram_counts = collections.Counter()
+    for order in range(1, max_order + 1):
+        for i in range(0, len(segment) - order + 1):
+            ngram = tuple(segment[i : i + order])
+            ngram_counts[ngram] += 1
+    return ngram_counts
+
+
+def compute_bleu(reference_corpus, translation_corpus, max_order=4, smooth=False):
+    """Computes BLEU score of translated segments against one or more references.
+
+    Args:
+      reference_corpus: list of lists of references for each translation. Each
+          reference should be tokenized into a list of tokens.
+      translation_corpus: list of translations to score. Each translation
+          should be tokenized into a list of tokens.
+      max_order: Maximum n-gram order to use when computing BLEU score.
+      smooth: Whether or not to apply Lin et al. 2004 smoothing.
+
+    Returns:
+      3-Tuple with the BLEU score, n-gram precisions, geometric mean of n-gram
+      precisions and brevity penalty.
+    """
+    matches_by_order = [0] * max_order
+    possible_matches_by_order = [0] * max_order
+    reference_length = 0
+    translation_length = 0
+    for references, translation in zip(reference_corpus, translation_corpus):
+        reference_length += min(len(r) for r in references)
+        translation_length += len(translation)
+
+        merged_ref_ngram_counts = collections.Counter()
+        for reference in references:
+            merged_ref_ngram_counts |= _get_ngrams(reference, max_order)
+        translation_ngram_counts = _get_ngrams(translation, max_order)
+        overlap = translation_ngram_counts & merged_ref_ngram_counts
+        for ngram in overlap:
+            matches_by_order[len(ngram) - 1] += overlap[ngram]
+        for order in range(1, max_order + 1):
+            possible_matches = len(translation) - order + 1
+            if possible_matches > 0:
+                possible_matches_by_order[order - 1] += possible_matches
+
+    precisions = [0] * max_order
+    for i in range(0, max_order):
+        if smooth:
+            precisions[i] = (matches_by_order[i] + 1.0) / (possible_matches_by_order[i] + 1.0)
+        else:
+            if possible_matches_by_order[i] > 0:
+                precisions[i] = float(matches_by_order[i]) / possible_matches_by_order[i]
+            else:
+                precisions[i] = 0.0
+
+    if min(precisions) > 0:
+        p_log_sum = sum((1.0 / max_order) * math.log(p) for p in precisions)
+        geo_mean = math.exp(p_log_sum)
+    else:
+        geo_mean = 0
+
+    ratio = float(translation_length) / reference_length
+
+    if ratio > 1.0:
+        bp = 1.0
+    else:
+        bp = math.exp(1 - 1.0 / ratio)
+
+    bleu = geo_mean * bp
+
+    return (bleu, precisions, bp, ratio, translation_length, reference_length)