chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,16 @@
|
||||
from FlagEmbedding.abc.evaluation import (
|
||||
AbsEvalArgs as MKQAEvalArgs,
|
||||
AbsEvalModelArgs as MKQAEvalModelArgs,
|
||||
)
|
||||
|
||||
from .data_loader import MKQAEvalDataLoader
|
||||
from .evaluator import MKQAEvaluator
|
||||
from .runner import MKQAEvalRunner
|
||||
|
||||
__all__ = [
|
||||
"MKQAEvalArgs",
|
||||
"MKQAEvalModelArgs",
|
||||
"MKQAEvalRunner",
|
||||
"MKQAEvalDataLoader",
|
||||
"MKQAEvaluator"
|
||||
]
|
||||
@@ -0,0 +1,28 @@
|
||||
from transformers import HfArgumentParser
|
||||
|
||||
from FlagEmbedding.evaluation.mkqa import (
|
||||
MKQAEvalArgs, MKQAEvalModelArgs,
|
||||
MKQAEvalRunner
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
parser = HfArgumentParser((
|
||||
MKQAEvalArgs,
|
||||
MKQAEvalModelArgs
|
||||
))
|
||||
|
||||
eval_args, model_args = parser.parse_args_into_dataclasses()
|
||||
eval_args: MKQAEvalArgs
|
||||
model_args: MKQAEvalModelArgs
|
||||
|
||||
runner = MKQAEvalRunner(
|
||||
eval_args=eval_args,
|
||||
model_args=model_args
|
||||
)
|
||||
|
||||
runner.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,228 @@
|
||||
import os
|
||||
import json
|
||||
import logging
|
||||
import datasets
|
||||
from tqdm import tqdm
|
||||
from typing import List, Optional
|
||||
|
||||
from FlagEmbedding.abc.evaluation import AbsEvalDataLoader
|
||||
|
||||
from .utils.normalize_text import normalize_text
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MKQAEvalDataLoader(AbsEvalDataLoader):
|
||||
"""
|
||||
Data loader class for MKQA.
|
||||
"""
|
||||
def available_dataset_names(self) -> List[str]:
|
||||
"""
|
||||
Get the available dataset names.
|
||||
|
||||
Returns:
|
||||
List[str]: All the available dataset names.
|
||||
"""
|
||||
return ['en', 'ar', 'fi', 'ja', 'ko', 'ru', 'es', 'sv', 'he', 'th', 'da', 'de', 'fr', 'it', 'nl', 'pl', 'pt', 'hu', 'vi', 'ms', 'km', 'no', 'tr', 'zh_cn', 'zh_hk', 'zh_tw']
|
||||
|
||||
def available_splits(self, dataset_name: Optional[str] = None) -> List[str]:
|
||||
"""
|
||||
Get the avaialble splits.
|
||||
|
||||
Args:
|
||||
dataset_name (str): Dataset name.
|
||||
|
||||
Returns:
|
||||
List[str]: All the available splits for the dataset.
|
||||
"""
|
||||
return ["test"]
|
||||
|
||||
def load_corpus(self, dataset_name: Optional[str] = None) -> datasets.DatasetDict:
|
||||
"""Load the corpus.
|
||||
|
||||
Args:
|
||||
dataset_name (Optional[str], optional): Name of the dataset. Defaults to None.
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: Loaded datasets instance of corpus.
|
||||
"""
|
||||
if self.dataset_dir is not None:
|
||||
# same corpus for all languages
|
||||
save_dir = self.dataset_dir
|
||||
return self._load_local_corpus(save_dir, dataset_name=dataset_name)
|
||||
else:
|
||||
return self._load_remote_corpus(dataset_name=dataset_name)
|
||||
|
||||
def _load_local_qrels(self, save_dir: str, dataset_name: Optional[str] = None, split: str = 'test') -> datasets.DatasetDict:
|
||||
"""Try to load qrels from local datasets.
|
||||
|
||||
Args:
|
||||
save_dir (str): Directory that save the data files.
|
||||
dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``.
|
||||
split (str, optional): Split of the dataset. Defaults to ``'test'``.
|
||||
|
||||
Raises:
|
||||
ValueError: No local qrels found, will try to download from remote.
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: Loaded datasets instance of qrels.
|
||||
"""
|
||||
checked_split = self.check_splits(split)
|
||||
if len(checked_split) == 0:
|
||||
raise ValueError(f"Split {split} not found in the dataset.")
|
||||
split = checked_split[0]
|
||||
|
||||
qrels_path = os.path.join(save_dir, f"{split}_qrels.jsonl")
|
||||
if self.force_redownload or not os.path.exists(qrels_path):
|
||||
logger.warning(f"Qrels not found in {qrels_path}. Trying to download the qrels from the remote and save it to {save_dir}.")
|
||||
return self._load_remote_qrels(dataset_name=dataset_name, split=split, save_dir=save_dir)
|
||||
else:
|
||||
qrels_data = datasets.load_dataset('json', data_files=qrels_path, cache_dir=self.cache_dir)['train']
|
||||
|
||||
qrels = {}
|
||||
for data in qrels_data:
|
||||
qid = data['qid']
|
||||
qrels[qid] = data['answers']
|
||||
|
||||
return datasets.DatasetDict(qrels)
|
||||
|
||||
def _load_remote_corpus(
|
||||
self,
|
||||
dataset_name: Optional[str] = None,
|
||||
save_dir: Optional[str] = None
|
||||
) -> datasets.DatasetDict:
|
||||
"""
|
||||
Refer to: https://arxiv.org/pdf/2402.03216. We use the corpus from the BeIR dataset.
|
||||
"""
|
||||
corpus = datasets.load_dataset(
|
||||
"BeIR/nq", "corpus",
|
||||
cache_dir=self.cache_dir,
|
||||
trust_remote_code=True,
|
||||
download_mode=self.hf_download_mode
|
||||
)["corpus"]
|
||||
|
||||
if save_dir is not None:
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
save_path = os.path.join(save_dir, "corpus.jsonl")
|
||||
corpus_dict = {}
|
||||
with open(save_path, "w", encoding="utf-8") as f:
|
||||
for data in tqdm(corpus, desc="Loading and Saving corpus"):
|
||||
docid, title, text = str(data["_id"]), normalize_text(data["title"]).lower(), normalize_text(data["text"]).lower()
|
||||
_data = {
|
||||
"id": docid,
|
||||
"title": title,
|
||||
"text": text
|
||||
}
|
||||
corpus_dict[docid] = {
|
||||
"title": title,
|
||||
"text": text
|
||||
}
|
||||
f.write(json.dumps(_data, ensure_ascii=False) + "\n")
|
||||
logging.info(f"{self.eval_name} corpus saved to {save_path}")
|
||||
else:
|
||||
corpus_dict = {}
|
||||
for data in tqdm(corpus, desc="Loading corpus"):
|
||||
docid, title, text = str(data["_id"]), normalize_text(data["title"]), normalize_text(data["text"])
|
||||
corpus_dict[docid] = {
|
||||
"title": title,
|
||||
"text": text
|
||||
}
|
||||
return datasets.DatasetDict(corpus_dict)
|
||||
|
||||
def _load_remote_qrels(
|
||||
self,
|
||||
dataset_name: str,
|
||||
split: str = 'test',
|
||||
save_dir: Optional[str] = None
|
||||
) -> datasets.DatasetDict:
|
||||
"""Load remote qrels from HF.
|
||||
|
||||
Args:
|
||||
dataset_name (str): Name of the dataset.
|
||||
split (str, optional): Split of the dataset. Defaults to ``'test'``.
|
||||
save_dir (Optional[str], optional): Directory to save the dataset. Defaults to ``None``.
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: Loaded datasets instance of qrel.
|
||||
"""
|
||||
endpoint = f"{os.getenv('HF_ENDPOINT', 'https://huggingface.co')}/datasets/Shitao/bge-m3-data"
|
||||
queries_download_url = f"{endpoint}/resolve/main/MKQA_test-data.zip"
|
||||
|
||||
qrels_save_dir = self._download_zip_file(queries_download_url, self.cache_dir)
|
||||
qrels_save_path = os.path.join(qrels_save_dir, f"{dataset_name}.jsonl")
|
||||
|
||||
if save_dir is not None:
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
save_path = os.path.join(save_dir, f"{split}_qrels.jsonl")
|
||||
qrels_dict = {}
|
||||
with open(save_path, "w", encoding="utf-8") as f1:
|
||||
with open(qrels_save_path, "r", encoding="utf-8") as f2:
|
||||
for line in tqdm(f2.readlines(), desc="Loading and Saving qrels"):
|
||||
data = json.loads(line)
|
||||
qid, answers = str(data["id"]), data["answers"]
|
||||
_data = {
|
||||
"qid": qid,
|
||||
"answers": answers
|
||||
}
|
||||
if qid not in qrels_dict:
|
||||
qrels_dict[qid] = {}
|
||||
qrels_dict[qid] = answers
|
||||
f1.write(json.dumps(_data, ensure_ascii=False) + "\n")
|
||||
logging.info(f"{self.eval_name} {dataset_name} qrels saved to {save_path}")
|
||||
else:
|
||||
qrels_dict = {}
|
||||
with open(qrels_save_path, "r", encoding="utf-8") as f:
|
||||
for line in tqdm(f.readlines(), desc="Loading qrels"):
|
||||
data = json.loads(line)
|
||||
qid, answers = str(data["id"]), data["answers"]
|
||||
if qid not in qrels_dict:
|
||||
qrels_dict[qid] = {}
|
||||
qrels_dict[qid] = answers
|
||||
return datasets.DatasetDict(qrels_dict)
|
||||
|
||||
def _load_remote_queries(
|
||||
self,
|
||||
dataset_name: str,
|
||||
split: str = 'test',
|
||||
save_dir: Optional[str] = None
|
||||
) -> datasets.DatasetDict:
|
||||
"""Load the queries from HF.
|
||||
|
||||
Args:
|
||||
dataset_name (str): Name of the dataset.
|
||||
split (str, optional): Split of the dataset. Defaults to ``'test'``.
|
||||
save_dir (Optional[str], optional): Directory to save the dataset. Defaults to ``None``.
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: Loaded datasets instance of queries.
|
||||
"""
|
||||
endpoint = f"{os.getenv('HF_ENDPOINT', 'https://huggingface.co')}/datasets/Shitao/bge-m3-data"
|
||||
queries_download_url = f"{endpoint}/resolve/main/MKQA_test-data.zip"
|
||||
|
||||
queries_save_dir = self._download_zip_file(queries_download_url, self.cache_dir)
|
||||
queries_save_path = os.path.join(queries_save_dir, f"{dataset_name}.jsonl")
|
||||
|
||||
if save_dir is not None:
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
save_path = os.path.join(save_dir, f"{split}_queries.jsonl")
|
||||
queries_dict = {}
|
||||
with open(save_path, "w", encoding="utf-8") as f1:
|
||||
with open(queries_save_path, "r", encoding="utf-8") as f2:
|
||||
for line in tqdm(f2.readlines(), desc="Loading and Saving queries"):
|
||||
data = json.loads(line)
|
||||
qid, query = str(data["id"]), data["question"]
|
||||
_data = {
|
||||
"id": qid,
|
||||
"text": query
|
||||
}
|
||||
queries_dict[qid] = query
|
||||
f1.write(json.dumps(_data, ensure_ascii=False) + "\n")
|
||||
logging.info(f"{self.eval_name} {dataset_name} queries saved to {save_path}")
|
||||
else:
|
||||
queries_dict = {}
|
||||
with open(queries_save_path, "r", encoding="utf-8") as f:
|
||||
for line in tqdm(f.readlines(), desc="Loading queries"):
|
||||
data = json.loads(line)
|
||||
qid, query = str(data["id"]), data["question"]
|
||||
queries_dict[qid] = query
|
||||
return datasets.DatasetDict(queries_dict)
|
||||
@@ -0,0 +1,117 @@
|
||||
import os
|
||||
from tqdm import tqdm
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from FlagEmbedding.abc.evaluation import AbsEvaluator
|
||||
|
||||
from .utils.compute_metrics import evaluate_qa_recall
|
||||
|
||||
|
||||
class MKQAEvaluator(AbsEvaluator):
|
||||
"""
|
||||
The evaluator class of MKQA.
|
||||
"""
|
||||
def get_corpus_embd_save_dir(
|
||||
self,
|
||||
retriever_name: str,
|
||||
corpus_embd_save_dir: Optional[str] = None,
|
||||
dataset_name: Optional[str] = None
|
||||
):
|
||||
"""Get the directory to save the corpus embedding.
|
||||
|
||||
Args:
|
||||
retriever_name (str): Name of the retriever.
|
||||
corpus_embd_save_dir (Optional[str], optional): Directory to save the corpus embedding. Defaults to ``None``.
|
||||
dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``.
|
||||
|
||||
Returns:
|
||||
str: The final directory to save the corpus embedding.
|
||||
"""
|
||||
if corpus_embd_save_dir is not None:
|
||||
# Save the corpus embeddings in the same directory for all dataset_name
|
||||
corpus_embd_save_dir = os.path.join(corpus_embd_save_dir, retriever_name)
|
||||
return corpus_embd_save_dir
|
||||
|
||||
def evaluate_results(
|
||||
self,
|
||||
search_results_save_dir: str,
|
||||
k_values: List[int] = [1, 3, 5, 10, 100, 1000]
|
||||
):
|
||||
"""Compute the metrics and get the eval results.
|
||||
|
||||
Args:
|
||||
search_results_save_dir (str): Directory that saves the search results.
|
||||
k_values (List[int], optional): Cutoffs. Defaults to ``[1, 3, 5, 10, 100, 1000]``.
|
||||
|
||||
Returns:
|
||||
dict: The evaluation results.
|
||||
"""
|
||||
eval_results_dict = {}
|
||||
|
||||
corpus = self.data_loader.load_corpus()
|
||||
corpus_dict = {}
|
||||
for docid, data in tqdm(corpus.items(), desc="Loading corpus for evaluation"):
|
||||
title, text = data["title"], data["text"]
|
||||
corpus_dict[docid] = f"{title} {text}".strip()
|
||||
|
||||
for file in os.listdir(search_results_save_dir):
|
||||
if not file.endswith('.json'):
|
||||
continue
|
||||
|
||||
file_path = os.path.join(search_results_save_dir, file)
|
||||
data_info, search_results = self.load_search_results(file_path)
|
||||
|
||||
_eval_name = data_info['eval_name']
|
||||
assert _eval_name == self.eval_name, f'Mismatch eval_name: {_eval_name} vs {self.eval_name} in {file_path}'
|
||||
|
||||
split = data_info['split']
|
||||
dataset_name = data_info.get('dataset_name', None)
|
||||
qrels = self.data_loader.load_qrels(dataset_name=dataset_name, split=split)
|
||||
|
||||
eval_results = self.compute_metrics(
|
||||
corpus_dict=corpus_dict,
|
||||
qrels=qrels,
|
||||
search_results=search_results,
|
||||
k_values=k_values
|
||||
)
|
||||
|
||||
if dataset_name is not None:
|
||||
key = f"{dataset_name}-{split}"
|
||||
else:
|
||||
key = split
|
||||
eval_results_dict[key] = eval_results
|
||||
|
||||
return eval_results_dict
|
||||
|
||||
@staticmethod
|
||||
def compute_metrics(
|
||||
corpus_dict: Dict[str, str],
|
||||
qrels: Dict[str, List[str]],
|
||||
search_results: Dict[str, Dict[str, float]],
|
||||
k_values: List[int],
|
||||
):
|
||||
"""
|
||||
Compute Recall@k for QA task. The definition of recall in QA task is different from the one in IR task. Please refer to the paper of RocketQA: https://aclanthology.org/2021.naacl-main.466.pdf.
|
||||
|
||||
Args:
|
||||
corpus_dict (Dict[str, str]): Dictionary of the corpus with doc id and contents.
|
||||
qrels (Dict[str, List[str]]): Relevances of queries and passage.
|
||||
search_results (Dict[str, Dict[str, float]]): Search results of the model to evaluate.
|
||||
|
||||
Returns:
|
||||
dict: The model's scores of the metrics.
|
||||
"""
|
||||
contexts = []
|
||||
answers = []
|
||||
top_k = max(k_values)
|
||||
for qid, doc_score_dict in search_results.items():
|
||||
doc_score_pair = sorted(doc_score_dict.items(), key=lambda x: x[1], reverse=True)
|
||||
_ctxs = [corpus_dict[docid] for docid, _ in doc_score_pair[:top_k]]
|
||||
contexts.append(_ctxs)
|
||||
answers.append(qrels[qid])
|
||||
|
||||
recall = evaluate_qa_recall(contexts, answers, k_values=k_values)
|
||||
scores = {f"qa_recall_at_{k}": v for k, v in zip(k_values, recall)}
|
||||
|
||||
return scores
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
from FlagEmbedding.abc.evaluation import AbsEvalRunner
|
||||
|
||||
from .data_loader import MKQAEvalDataLoader
|
||||
from .evaluator import MKQAEvaluator
|
||||
|
||||
|
||||
class MKQAEvalRunner(AbsEvalRunner):
|
||||
"""
|
||||
Evaluation runner of MKQA.
|
||||
"""
|
||||
def load_data_loader(self) -> MKQAEvalDataLoader:
|
||||
"""Load the data loader instance by args.
|
||||
|
||||
Returns:
|
||||
MKQAEvalDataLoader: The MKQA data loader instance.
|
||||
"""
|
||||
data_loader = MKQAEvalDataLoader(
|
||||
eval_name=self.eval_args.eval_name,
|
||||
dataset_dir=self.eval_args.dataset_dir,
|
||||
cache_dir=self.eval_args.cache_path,
|
||||
token=self.eval_args.token,
|
||||
force_redownload=self.eval_args.force_redownload,
|
||||
)
|
||||
return data_loader
|
||||
|
||||
def load_evaluator(self) -> MKQAEvaluator:
|
||||
"""Load the evaluator instance by args.
|
||||
|
||||
Returns:
|
||||
MKQAEvaluator: The MKQA evaluator instance.
|
||||
"""
|
||||
evaluator = MKQAEvaluator(
|
||||
eval_name=self.eval_args.eval_name,
|
||||
data_loader=self.data_loader,
|
||||
overwrite=self.eval_args.overwrite,
|
||||
)
|
||||
return evaluator
|
||||
@@ -0,0 +1,95 @@
|
||||
"""
|
||||
Ref: https://github.com/facebookresearch/contriever
|
||||
"""
|
||||
import regex
|
||||
import unicodedata
|
||||
from functools import partial
|
||||
from typing import List, Union
|
||||
|
||||
|
||||
class SimpleTokenizer:
|
||||
ALPHA_NUM = r'[\p{L}\p{N}\p{M}]+'
|
||||
NON_WS = r'[^\p{Z}\p{C}]'
|
||||
|
||||
def __init__(self):
|
||||
"""
|
||||
Args:
|
||||
annotators: None or empty set (only tokenizes).
|
||||
"""
|
||||
self._regexp = regex.compile(
|
||||
'(%s)|(%s)' % (self.ALPHA_NUM, self.NON_WS),
|
||||
flags=regex.IGNORECASE + regex.UNICODE + regex.MULTILINE
|
||||
)
|
||||
|
||||
def tokenize(self, text, uncased=False):
|
||||
matches = [m for m in self._regexp.finditer(text)]
|
||||
if uncased:
|
||||
tokens = [m.group().lower() for m in matches]
|
||||
else:
|
||||
tokens = [m.group() for m in matches]
|
||||
return tokens
|
||||
|
||||
|
||||
def _normalize(text):
|
||||
return unicodedata.normalize('NFD', text)
|
||||
|
||||
|
||||
def has_answer(answers, text, tokenizer) -> bool:
|
||||
"""Check if a document contains an answer string."""
|
||||
text = _normalize(text)
|
||||
text = tokenizer.tokenize(text, uncased=True)
|
||||
|
||||
for answer in answers:
|
||||
answer = _normalize(answer)
|
||||
answer = tokenizer.tokenize(answer, uncased=True)
|
||||
for i in range(0, len(text) - len(answer) + 1):
|
||||
if answer == text[i: i + len(answer)]:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def check_answer(example, tokenizer) -> List[bool]:
|
||||
"""Search through all the top docs to see if they have any of the answers."""
|
||||
answers = example['answers']
|
||||
ctxs = example['ctxs']
|
||||
|
||||
hits = []
|
||||
for i, text in enumerate(ctxs):
|
||||
if text is None: # cannot find the document for some reason
|
||||
hits.append(False)
|
||||
continue
|
||||
hits.append(has_answer(answers, text, tokenizer))
|
||||
return hits
|
||||
|
||||
|
||||
def evaluate_qa_recall(ctxs, answers, k_values: Union[int, List[int]]=100):
|
||||
# compute Recall@k for QA task
|
||||
data = []
|
||||
assert len(ctxs) == len(answers)
|
||||
for i in range(len(ctxs)):
|
||||
_ctxs, _answers = ctxs[i], answers[i]
|
||||
data.append({
|
||||
'answers': _answers,
|
||||
'ctxs': _ctxs,
|
||||
})
|
||||
tokenizer = SimpleTokenizer()
|
||||
get_score_partial = partial(check_answer, tokenizer=tokenizer)
|
||||
|
||||
scores = map(get_score_partial, data)
|
||||
|
||||
n_docs = len(data[0]['ctxs'])
|
||||
top_k_hits = [0] * n_docs
|
||||
for question_hits in scores:
|
||||
best_hit = next((i for i, x in enumerate(question_hits) if x), None)
|
||||
if best_hit is not None:
|
||||
top_k_hits[best_hit:] = [v + 1 for v in top_k_hits[best_hit:]]
|
||||
|
||||
if isinstance(k_values, int):
|
||||
k = min(k_values, len(top_k_hits))
|
||||
return top_k_hits[k - 1] / len(data)
|
||||
else:
|
||||
scores = []
|
||||
for k in k_values:
|
||||
k = min(k, len(top_k_hits))
|
||||
scores.append(top_k_hits[k - 1] / len(data))
|
||||
return scores
|
||||
@@ -0,0 +1,162 @@
|
||||
"""
|
||||
adapted from chemdataextractor.text.normalize
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
Tools for normalizing text.
|
||||
https://github.com/mcs07/ChemDataExtractor
|
||||
:copyright: Copyright 2016 by Matt Swain.
|
||||
:license: MIT
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining
|
||||
a copy of this software and associated documentation files (the
|
||||
'Software'), to deal in the Software without restriction, including
|
||||
without limitation the rights to use, copy, modify, merge, publish,
|
||||
distribute, sublicense, and/or sell copies of the Software, and to
|
||||
permit persons to whom the Software is furnished to do so, subject to
|
||||
the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be
|
||||
included in all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED 'AS IS', WITHOUT WARRANTY OF ANY KIND,
|
||||
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
|
||||
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
|
||||
IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
|
||||
CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
|
||||
TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
|
||||
SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
||||
"""
|
||||
|
||||
#: Control characters.
|
||||
CONTROLS = {
|
||||
'\u0001', '\u0002', '\u0003', '\u0004', '\u0005', '\u0006', '\u0007', '\u0008', '\u000e', '\u000f', '\u0011',
|
||||
'\u0012', '\u0013', '\u0014', '\u0015', '\u0016', '\u0017', '\u0018', '\u0019', '\u001a', '\u001b',
|
||||
}
|
||||
# There are further control characters, but they are instead replaced with a space by unicode normalization
|
||||
# '\u0009', '\u000a', '\u000b', '\u000c', '\u000d', '\u001c', '\u001d', '\u001e', '\u001f'
|
||||
|
||||
|
||||
#: Hyphen and dash characters.
|
||||
HYPHENS = {
|
||||
'-', # \u002d Hyphen-minus
|
||||
'‐', # \u2010 Hyphen
|
||||
'‑', # \u2011 Non-breaking hyphen
|
||||
'⁃', # \u2043 Hyphen bullet
|
||||
'‒', # \u2012 figure dash
|
||||
'–', # \u2013 en dash
|
||||
'—', # \u2014 em dash
|
||||
'―', # \u2015 horizontal bar
|
||||
}
|
||||
|
||||
#: Minus characters.
|
||||
MINUSES = {
|
||||
'-', # \u002d Hyphen-minus
|
||||
'−', # \u2212 Minus
|
||||
'－', # \uff0d Full-width Hyphen-minus
|
||||
'⁻', # \u207b Superscript minus
|
||||
}
|
||||
|
||||
#: Plus characters.
|
||||
PLUSES = {
|
||||
'+', # \u002b Plus
|
||||
'＋', # \uff0b Full-width Plus
|
||||
'⁺', # \u207a Superscript plus
|
||||
}
|
||||
|
||||
#: Slash characters.
|
||||
SLASHES = {
|
||||
'/', # \u002f Solidus
|
||||
'⁄', # \u2044 Fraction slash
|
||||
'∕', # \u2215 Division slash
|
||||
}
|
||||
|
||||
#: Tilde characters.
|
||||
TILDES = {
|
||||
'~', # \u007e Tilde
|
||||
'˜', # \u02dc Small tilde
|
||||
'⁓', # \u2053 Swung dash
|
||||
'∼', # \u223c Tilde operator #in mbert vocab
|
||||
'∽', # \u223d Reversed tilde
|
||||
'∿', # \u223f Sine wave
|
||||
'〜', # \u301c Wave dash #in mbert vocab
|
||||
'～', # \uff5e Full-width tilde #in mbert vocab
|
||||
}
|
||||
|
||||
#: Apostrophe characters.
|
||||
APOSTROPHES = {
|
||||
"'", # \u0027
|
||||
'’', # \u2019
|
||||
'՚', # \u055a
|
||||
'Ꞌ', # \ua78b
|
||||
'ꞌ', # \ua78c
|
||||
'＇', # \uff07
|
||||
}
|
||||
|
||||
#: Single quote characters.
|
||||
SINGLE_QUOTES = {
|
||||
"'", # \u0027
|
||||
'‘', # \u2018
|
||||
'’', # \u2019
|
||||
'‚', # \u201a
|
||||
'‛', # \u201b
|
||||
|
||||
}
|
||||
|
||||
#: Double quote characters.
|
||||
DOUBLE_QUOTES = {
|
||||
'"', # \u0022
|
||||
'“', # \u201c
|
||||
'”', # \u201d
|
||||
'„', # \u201e
|
||||
'‟', # \u201f
|
||||
}
|
||||
|
||||
#: Accent characters.
|
||||
ACCENTS = {
|
||||
'`', # \u0060
|
||||
'´', # \u00b4
|
||||
}
|
||||
|
||||
#: Prime characters.
|
||||
PRIMES = {
|
||||
'′', # \u2032
|
||||
'″', # \u2033
|
||||
'‴', # \u2034
|
||||
'‵', # \u2035
|
||||
'‶', # \u2036
|
||||
'‷', # \u2037
|
||||
'⁗', # \u2057
|
||||
}
|
||||
|
||||
#: Quote characters, including apostrophes, single quotes, double quotes, accents and primes.
|
||||
QUOTES = APOSTROPHES | SINGLE_QUOTES | DOUBLE_QUOTES | ACCENTS | PRIMES
|
||||
|
||||
def normalize_text(text: str):
|
||||
for control in CONTROLS:
|
||||
text = text.replace(control, '')
|
||||
text = text.replace('\u000b', ' ').replace('\u000c', ' ').replace(u'\u0085', ' ')
|
||||
|
||||
for hyphen in HYPHENS | MINUSES:
|
||||
text = text.replace(hyphen, '-')
|
||||
text = text.replace('\u00ad', '')
|
||||
|
||||
for double_quote in DOUBLE_QUOTES:
|
||||
text = text.replace(double_quote, '"') # \u0022
|
||||
for single_quote in (SINGLE_QUOTES | APOSTROPHES | ACCENTS):
|
||||
text = text.replace(single_quote, "'") # \u0027
|
||||
text = text.replace('′', "'") # \u2032 prime
|
||||
text = text.replace('‵', "'") # \u2035 reversed prime
|
||||
text = text.replace('″', "''") # \u2033 double prime
|
||||
text = text.replace('‶', "''") # \u2036 reversed double prime
|
||||
text = text.replace('‴', "'''") # \u2034 triple prime
|
||||
text = text.replace('‷', "'''") # \u2037 reversed triple prime
|
||||
text = text.replace('⁗', "''''") # \u2057 quadruple prime
|
||||
|
||||
text = text.replace('…', '...').replace(' . . . ', ' ... ') # \u2026
|
||||
|
||||
for slash in SLASHES:
|
||||
text = text.replace(slash, '/')
|
||||
|
||||
#for tilde in TILDES:
|
||||
# text = text.replace(tilde, '~')
|
||||
|
||||
return text
|
||||
Reference in New Issue
Block a user