def test_expand_vocabulary(self):
generator = CandidatesGenerator()
generator.construct(VOCABULARY_FILE, VOCABULARY_FILE,
FASTTEXT_DUMP_FILE, self.config)
additional_tokens = {"a", "aaa", "123", "get", "341"}
vocabulary = generator.tokens
generator.expand_vocabulary(additional_tokens)
self.assertSetEqual(generator.tokens,
vocabulary.union(additional_tokens))
def test_expand_vocabulary(self):
generator = CandidatesGenerator()
generator.construct(VOCABULARY_FILE,
VOCABULARY_FILE,
FASTTEXT_DUMP_FILE,
neighbors=3,
edit_candidates=3,
max_distance=3,
radius=3)
additional_tokens = {"a", "aaa", "123", "get", "341"}
vocabulary = generator.tokens
generator.expand_vocabulary(additional_tokens)
self.assertSetEqual(generator.tokens,
vocabulary.union(additional_tokens))
class TyposCorrector(Model):
"""
Model for correcting typos in tokens inside identifiers.
"""
_log = logging.getLogger("TyposCorrector")
NAME = "typos_correction"
VENDOR = "source{d}"
DESCRIPTION = "Model that suggests fixes to correct typos."
LICENSE = DEFAULT_LICENSE
def __init__(self, ranking_config: Optional[Mapping[str, Any]] = None, **kwargs):
"""
Initialize a new instance of TyposCorrector class.
:param ranking_config: Ranking configuration, options:
train_rounds: Number of training rounds (int).
early_stopping: Early stopping parameter (int).
boost_param: Boosting parameters (dict).
:param kwargs: Extra keyword arguments which are consumed by Model.
"""
super().__init__(**kwargs)
self.generator = CandidatesGenerator()
self.ranker = CandidatesRanker(ranking_config)
@property
def processes_number(self) -> int:
"""Return the number of processes for multiprocessing used to train and to predict."""
return self.ranker.config["boost_param"]["nthread"]
@processes_number.setter
def processes_number(self, processes_number: int):
"""Set the number of processes for multiprocessing used to train and to predict."""
self.ranker.config["boost_param"]["nthread"] = processes_number
def initialize_generator(self, vocabulary_file: str, frequencies_file: str,
embeddings_file: str, config: Optional[Mapping[str, Any]] = None,
) -> None:
"""
Construct a new CandidatesGenerator.
:param vocabulary_file: The path to the vocabulary.
:param frequencies_file: The path to the frequencies.
:param embeddings_file: The path to the embeddings.
:param config: Candidates generation configuration, options:
neighbors_number: Number of neighbors of context and typo embeddings \
to consider as candidates (int).
edit_dist_number: Number of the most frequent tokens among tokens at \
equal edit distance from the typo to consider as \
candidates (int).
max_distance: Maximum edit distance for symspell lookup for candidates \
(int).
radius: Maximum edit distance from typo allowed for candidates (int).
max_corrected_length: Maximum length of prefix in which symspell lookup \
for typos is conducted (int).
start_pool_size: Length of data, starting from which multiprocessing is \
desired (int).
chunksize: Max size of a chunk for one process during multiprocessing (int).
"""
self.generator.construct(vocabulary_file, frequencies_file, embeddings_file, config)
self._log.debug("%s is initialized", repr(self.generator))
def set_ranking_config(self, config: Mapping[str, Any]) -> None:
"""
Update the ranking config - see XGBoost docs for details.
:param config: Ranking configuration, options:
train_rounds: Number of training rounds (int).
early_stopping: Early stopping parameter (int).
boost_param: Boosting parameters (dict).
"""
self.ranker.set_config(config)
self._log.debug("%s is initialized", repr(self.ranker))
def set_generation_config(self, config: Mapping[str, Any]) -> None:
"""
Update the candidates generation config.
:param config: Candidates generation configuration, options:
neighbors_number: Number of neighbors of context and typo embeddings \
to consider as candidates (int).
edit_dist_number: Number of the most frequent tokens among tokens at \
equal edit distance from the typo to consider as \
candidates (int).
max_distance: Maximum edit distance for symspell lookup for candidates \
(int).
radius: Maximum edit distance from typo allowed for candidates (int).
max_corrected_length: Maximum length of prefix in which symspell lookup \
for typos is conducted (int).
start_pool_size: Length of data, starting from which multiprocessing is \
desired (int).
chunksize: Max size of a chunk for one process during multiprocessing (int).
"""
self.generator.set_config(config)
self._log.debug("%s is initialized", repr(self.ranker))
def expand_vocabulary(self, additional_tokens: Iterable[str]) -> None:
"""
Add given tokens to the model's vocabulary.
:param additional_tokens: Tokens to add to the vocabulary.
"""
self.generator.expand_vocabulary(additional_tokens)
def train(self, data: pandas.DataFrame, candidates: Optional[str] = None,
save_candidates_file: Optional[str] = None) -> None:
"""
Train corrector on tokens from the given dataset.
:param data: DataFrame which contains columns Columns.Token, Columns.CorrectToken, \
and Columns.Split.
:param candidates: A .csv.xz dump of a dataframe with precalculated candidates.
:param save_candidates_file: Path to file where to save the candidates (.csv.xz).
"""
self._log.info("train input shape: %s", data.shape)
if candidates is None:
self._log.info("candidates were not provided and will be generated")
candidates = self.generator.generate_candidates(
data, self.processes_number, save_candidates_file)
else:
candidates = pandas.read_csv(candidates, index_col=0, keep_default_na=False)
self._log.info("loaded candidates from %s", candidates)
self.ranker.fit(data[Columns.CorrectToken], get_candidates_metadata(candidates),
get_candidates_features(candidates))
def train_on_file(self, data_file: str, candidates: Optional[str] = None,
save_candidates_file: Optional[str] = None) -> None:
"""
Train corrector on tokens from the given file.
:param data_file: A .csv dump of a dataframe which contains columns Columns.Token, \
Columns.CorrectToken and Columns.Split.
:param candidates: A .csv.xz dump of a dataframe with precalculated candidates.
:param save_candidates_file: Path to file where to save the candidates (.csv.xz).
"""
self.train(pandas.read_csv(data_file, index_col=0, keep_default_na=False), candidates,
save_candidates_file)
def suggest(self, data: pandas.DataFrame, candidates: Optional[str] = None,
save_candidates_file: Optional[str] = None, n_candidates: int = 3,
return_all: bool = True) -> Dict[int, List[Candidate]]:
"""
Suggest corrections for the tokens from the given dataset.
:param data: DataFrame which contains columns Columns.Token and Columns.Split.
:param candidates: A .csv.xz dump of a dataframe with precalculated candidates.
:param save_candidates_file: Path to file to save candidates to (.csv.xz).
:param n_candidates: Number of most probable candidates to return.
:param return_all: False to return suggestions only for corrected tokens.
:return: Dictionary `{id : [(candidate, correctness_proba), ...]}`, candidates are sorted \
by correctness probability in a descending order.
"""
if candidates is None:
candidates = self.generator.generate_candidates(
data, self.processes_number, save_candidates_file)
else:
candidates = pandas.read_csv(candidates, index_col=0, keep_default_na=False)
return self.ranker.rank(get_candidates_metadata(candidates),
get_candidates_features(candidates), n_candidates, return_all)
def suggest_on_file(self, data_file: str, candidates: Optional[str] = None,
save_candidates_file: Optional[str] = None, n_candidates: int = 3,
return_all: bool = True) -> Dict[int, List[Candidate]]:
"""
Suggest corrections for the tokens from the given file.
:param data_file: A .csv dump of a DataFrame which contains columns Columns.Token \
and Columns.Split.
:param candidates: A .csv.xz dump of a dataframe with precalculated candidates.
:param save_candidates_file: Path to file to save candidates to (.csv.xz).
:param n_candidates: Number of most probable candidates to return.
:param return_all: False to return suggestions only for corrected tokens.
:return: Dictionary `{id : [(candidate, correctness_proba), ...]}`, candidates are sorted \
by correctness probability in a descending order.
"""
return self.suggest(pandas.read_csv(data_file, index_col=0, keep_default_na=False),
candidates, save_candidates_file, n_candidates, return_all)
def suggest_by_batches(self, data: pandas.DataFrame, n_candidates: int = 3,
return_all: bool = True, batch_size: int = 2048,
) -> Dict[int, List[Candidate]]:
"""
Suggest corrections for the tokens from the given dataset by batches. \
Does not support precalculated candidates.
:param data: DataFrame which contains columns Columns.Token and Columns.Split.
:param n_candidates: Number of most probable candidates to return.
:param return_all: False to return suggestions only for corrected tokens.
:param batch_size: Batch size.
:return: Dictionary `{id : [(candidate, correctness_proba), ...]}`, candidates are sorted \
by correctness probability in a descending order.
"""
all_suggestions = []
for i in tqdm(range(0, len(data), batch_size)):
suggestions = self.suggest(data.iloc[i:i + batch_size, :], n_candidates=n_candidates,
return_all=return_all)
all_suggestions.append(suggestions.items())
return dict(chain.from_iterable(all_suggestions))
def evaluate(self, test_data: pandas.DataFrame) -> Tuple[Dict[int, List[Candidate]], str]:
"""
Evaluate the corrector on the given test dataset.
Save the result metrics to the model metadata and print it to the standard output.
:param test_data: DataFrame which contains column Columns.Token, \
column Columns.Split is optional, but used when present.
:return: Suggestions for correction of tokens inside the `test_data` and the quality
report.
"""
self._log.info("evaluate on test data with shape %s", test_data.shape)
suggestions = self.suggest(test_data)
report = generate_report(test_data, suggestions)
self.metrics = get_scores(test_data, suggestions)
self._log.info("evaluation report:\n%s", report)
return suggestions, report
def __eq__(self, other: "TyposCorrector") -> bool:
return self.generator == other.generator and self.ranker == other.ranker
def dump(self) -> str:
"""Model.__str__ to format the object."""
return ("# Generator\n"
"%s\n\n"
"# Ranker\n"
"%s" %
(self.generator.dump(), self.ranker.dump()))
def _generate_tree(self) -> dict:
return {"generator": self.generator._generate_tree(),
"ranker": self.ranker._generate_tree()}
def _load_tree(self, tree: dict) -> None:
self.generator._load_tree(tree["generator"])
self.ranker._load_tree(tree["ranker"])
class TyposCorrector(Model):
"""
Model for correcting typos in tokens inside identifiers.
"""
_log = logging.getLogger("TyposCorrector")
NAME = "typos_correction"
VENDOR = "source{d}"
DESCRIPTION = "Model that suggests fixes to correct typos."
LICENSE = DEFAULT_LICENSE
DEFAULT_RADIUS = 3
DEFAULT_MAX_DISTANCE = 2
DEFAULT_NEIGHBORS_NUMBER = 0
DEFAULT_EDIT_CANDIDATES = 20
DEFAULT_TRAIN_ROUNDS = 4000
DEFAULT_EARLY_STOPPING = 200
DEFAULT_BOOST_PARAM = {"max_depth": 6,
"eta": 0.03,
"min_child_weight": 2,
"silent": 1,
"objective": "binary:logistic",
"subsample": 0.5,
"colsample_bytree": 0.5,
"alpha": 1,
"eval_metric": ["error"],
"nthread": 0}
DEFAULT_START_POOL_SIZE = 64
DEFAULT_CHUNKSIZE = 256
def __init__(self, **kwargs):
"""
Initialize a new instance of TyposCorrector class.
:param kwargs: extra keyword arguments which are consumed by Model.
"""
super().__init__(**kwargs)
self.generator = CandidatesGenerator()
self.ranker = CandidatesRanker()
@property
def processes_number(self) -> int:
"""Return the number of processes for multiprocessing used to train and to predict."""
return self.ranker.boost_param["nthread"]
@processes_number.setter
def processes_number(self, processes_number: int):
"""Set the number of processes for multiprocessing used to train and to predict."""
self.ranker.boost_param["nthread"] = processes_number
def initialize_ranker(self, boost_params: Optional[dict] = None,
train_rounds: int = DEFAULT_TRAIN_ROUNDS,
early_stopping: int = DEFAULT_EARLY_STOPPING) -> None:
"""
Apply the ranking parameters - see XGBoost docs for details.
:param train_rounds: Number of training rounds.
:param early_stopping: Early stopping parameter.
:param boost_params: Boosting parameters. The defaults are DEFAULT_BOOST_PARAM.
:return: Nothing
"""
boost_params = boost_params or self.DEFAULT_BOOST_PARAM
self.ranker.construct(boost_params, train_rounds, early_stopping)
self._log.debug("%s is initialized", repr(self.ranker))
def initialize_generator(self, vocabulary_file: str, frequencies_file: str,
embeddings_file: Optional[str] = None,
neighbors_number: int = DEFAULT_NEIGHBORS_NUMBER,
edit_candidates: int = DEFAULT_EDIT_CANDIDATES,
max_distance: int = DEFAULT_MAX_DISTANCE,
radius: int = DEFAULT_RADIUS) -> None:
"""
Construct a new CandidatesGenerator.
:param vocabulary_file: The path to the vocabulary.
:param frequencies_file: The path to the frequencies.
:param embeddings_file: The path to the embeddings.
:param neighbors_number: Number of neighbors of context and typo embeddings \
to consider as candidates.
:param edit_candidates: Number of the most frequent tokens among tokens on \
equal edit distance from the typo to consider as candidates.
:param max_distance: Maximum edit distance for symspell lookup.
:param radius: Maximum edit distance from typo allowed for candidates.
"""
self.generator.construct(vocabulary_file, frequencies_file, embeddings_file,
neighbors_number, edit_candidates, max_distance, radius)
self._log.debug("%s is initialized", repr(self.generator))
def expand_vocabulary(self, additional_tokens: Set[str]) -> None:
"""
Add given tokens to the model's vocabulary.
:param additional_tokens: Set of tokens to add to the vocabulary.
"""
self.generator.expand_vocabulary(additional_tokens)
def train(self, data: pandas.DataFrame, candidates: Optional[str] = None,
save_candidates_file: Optional[str] = None,
start_pool_size: int = DEFAULT_START_POOL_SIZE,
chunksize: int = DEFAULT_CHUNKSIZE) -> None:
"""
Train corrector on tokens from the given dataset.
:param data: DataFrame which contains columns Columns.Token, Columns.CorrectToken, \
and Columns.Split.
:param candidates: A .csv.xz dump of a dataframe with precalculated candidates.
:param save_candidates_file: Path to file where to save the candidates (.csv.xz).
:param start_pool_size: Length of data, starting from which multiprocessing is desired.
:param chunksize: Max size of a chunk for one process during multiprocessing.
"""
self._log.info("train input shape: %s", data.shape)
if candidates is None:
self._log.info("candidates were not provided and will be generated")
candidates = self.generator.generate_candidates(
data, self.processes_number, start_pool_size, chunksize, save_candidates_file)
else:
candidates = pandas.read_csv(candidates, index_col=0, keep_default_na=False)
self._log.info("loaded candidates from %s", candidates)
self.ranker.fit(data[Columns.CorrectToken], get_candidates_metadata(candidates),
get_candidates_features(candidates))
def train_on_file(self, data_file: str, candidates: Optional[str] = None,
save_candidates_file: Optional[str] = None,
start_pool_size: int = DEFAULT_START_POOL_SIZE,
chunksize: int = DEFAULT_CHUNKSIZE) -> None:
"""
Train corrector on tokens from the given file.
:param data_file: A .csv dump of a dataframe which contains columns Columns.Token, \
Columns.CorrectToken and Columns.Split.
:param candidates: A .csv.xz dump of a dataframe with precalculated candidates.
:param save_candidates_file: Path to file where to save the candidates (.csv.xz).
:param start_pool_size: Length of data, starting from which multiprocessing is desired.
:param chunksize: Max size of a chunk for one process during multiprocessing.
"""
self.train(pandas.read_csv(data_file, index_col=0, keep_default_na=False), candidates,
save_candidates_file, start_pool_size, chunksize)
def suggest(self, data: pandas.DataFrame, candidates: Optional[str] = None,
save_candidates_file: Optional[str] = None, n_candidates: int = 3,
return_all: bool = True, start_pool_size: int = DEFAULT_START_POOL_SIZE,
chunksize: int = DEFAULT_CHUNKSIZE) -> Dict[int, List[Tuple[str, float]]]:
"""
Suggest corrections for the tokens from the given dataset.
:param data: DataFrame which contains columns Columns.Token and Columns.Split.
:param candidates: A .csv.xz dump of a dataframe with precalculated candidates.
:param save_candidates_file: Path to file to save candidates to (.csv.xz).
:param n_candidates: Number of most probable candidates to return.
:param return_all: False to return suggestions only for corrected tokens.
:param start_pool_size: Length of data, starting from which multiprocessing is desired.
:param chunksize: Max size of a chunk for one process during multiprocessing.
:return: Dictionary `{id : [(candidate, correctness_proba), ...]}`, candidates are sorted \
by correctness probability in a descending order.
"""
if candidates is None:
candidates = self.generator.generate_candidates(
data, self.processes_number, start_pool_size, chunksize, save_candidates_file)
else:
candidates = pandas.read_csv(candidates, index_col=0, keep_default_na=False)
return self.ranker.rank(get_candidates_metadata(candidates),
get_candidates_features(candidates), n_candidates, return_all)
def suggest_on_file(self, data_file: str, candidates: Optional[str] = None,
save_candidates_file: Optional[str] = None, n_candidates: int = 3,
return_all: bool = True, start_pool_size: int = DEFAULT_START_POOL_SIZE,
chunksize: int = DEFAULT_CHUNKSIZE) -> Dict[int, List[Tuple[str, float]]]:
"""
Suggest corrections for the tokens from the given file.
:param data_file: A .csv dump of a DataFrame which contains columns Columns.Token \
and Columns.Split.
:param candidates: A .csv.xz dump of a dataframe with precalculated candidates.
:param save_candidates_file: Path to file to save candidates to (.csv.xz).
:param n_candidates: Number of most probable candidates to return.
:param return_all: False to return suggestions only for corrected tokens.
:param start_pool_size: Length of data, starting from which multiprocessing is desired.
:param chunksize: Max size of a chunk for one process during multiprocessing.
:return: Dictionary `{id : [(candidate, correctness_proba), ...]}`, candidates are sorted \
by correctness probability in a descending order.
"""
return self.suggest(pandas.read_csv(data_file, index_col=0, keep_default_na=False),
candidates, save_candidates_file, n_candidates, return_all,
start_pool_size, chunksize)
def suggest_by_batches(self, data: pandas.DataFrame, n_candidates: int = 3,
return_all: bool = True, batch_size: int = 2048,
start_pool_size: int = DEFAULT_START_POOL_SIZE,
chunksize: int = DEFAULT_CHUNKSIZE,
) -> Dict[int, List[Tuple[str, float]]]:
"""
Suggest corrections for the tokens from the given dataset by batches. \
Does not support precalculated candidates.
:param data: DataFrame which contains columns Columns.Token and Columns.Split.
:param n_candidates: Number of most probable candidates to return.
:param return_all: False to return suggestions only for corrected tokens.
:param batch_size: Batch size.
:param start_pool_size: Length of data, starting from which multiprocessing is desired.
:param chunksize: Max size of a chunk for one process during multiprocessing.
:return: Dictionary `{id : [(candidate, correctness_proba), ...]}`, candidates are sorted \
by correctness probability in a descending order.
"""
all_suggestions = []
for i in tqdm(range(0, len(data), batch_size)):
suggestions = self.suggest(data.iloc[i:i + batch_size, :], n_candidates=n_candidates,
return_all=return_all, start_pool_size=start_pool_size,
chunksize=chunksize)
all_suggestions.append(suggestions.items())
return dict(chain.from_iterable(all_suggestions))
def evaluate(self, test_data: pandas.DataFrame, start_pool_size: int = DEFAULT_START_POOL_SIZE,
chunksize: int = DEFAULT_CHUNKSIZE) -> None:
"""
Evaluate the corrector on the given test dataset.
Save the result metrics to the model metadata and print it to the standard output.
:param test_data: DataFrame which contains column Columns.Token, \
column Columns.Split is optional, but used when present
:param start_pool_size: Length of data, starting from which multiprocessing is desired.
:param chunksize: Max size of a chunk for one process during multiprocessing.
"""
self._log.info("evaluate on test data with shape %s", test_data.shape)
suggestions = self.suggest(test_data, start_pool_size=start_pool_size,
chunksize=chunksize)
self.metrics = get_scores(test_data, suggestions)
self._log.info("evaluation report:\n%s", generate_report(test_data, suggestions))
def __eq__(self, other: "TyposCorrector") -> bool:
return self.generator == other.generator and self.ranker == other.ranker
def dump(self) -> str:
"""Model.__str__ to format the object."""
return ("# Generator\n"
"%s\n\n"
"# Ranker\n"
"%s" %
(self.generator.dump(), self.ranker.dump()))
def _generate_tree(self) -> dict:
return {"generator": self.generator._generate_tree(),
"ranker": self.ranker._generate_tree()}
def _load_tree(self, tree: dict) -> None:
self.generator._load_tree(tree["generator"])
self.ranker._load_tree(tree["ranker"])