diff options
author | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2021-03-29 21:39:58 +0200 |
---|---|---|
committer | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2021-03-29 21:39:58 +0200 |
commit | d21594211e29c40c135b753e33b248b0737cd76f (patch) | |
tree | ea37e2701fdc45b27815d8831e6b60ff6888e168 /text_recognizer | |
parent | 46a1472d33d3a4180798492e819f2ec02bc3b1a3 (diff) |
Refactor word piece scripts
Diffstat (limited to 'text_recognizer')
-rw-r--r-- | text_recognizer/data/build_transitions.py | 263 | ||||
-rw-r--r-- | text_recognizer/data/iam_preprocessor.py | 2 | ||||
-rw-r--r-- | text_recognizer/data/make_wordpieces.py | 114 |
3 files changed, 378 insertions, 1 deletions
diff --git a/text_recognizer/data/build_transitions.py b/text_recognizer/data/build_transitions.py new file mode 100644 index 0000000..91f8c1a --- /dev/null +++ b/text_recognizer/data/build_transitions.py @@ -0,0 +1,263 @@ +"""Builds transition graph. + +Most code stolen from here: + + https://github.com/facebookresearch/gtn_applications/blob/master/scripts/build_transitions.py + +""" + +import collections +import itertools +from pathlib import Path +from typing import Any, Dict, List, Optional + +import click +import gtn +from loguru import logger + + +START_IDX = -1 +END_IDX = -2 +WORDSEP = "▁" + + +def build_graph(ngrams: List, disable_backoff: bool = False) -> gtn.Graph: + """Returns a gtn Graph based on the ngrams.""" + graph = gtn.Graph(False) + ngram = len(ngrams) + state_to_node = {} + + def get_node(state: Optional[List]) -> Any: + node = state_to_node.get(state, None) + + if node is not None: + return node + + start = state == tuple([START_IDX]) if ngram > 1 else True + end = state == tuple([END_IDX]) if ngram > 1 else True + node = graph.add_node(start, end) + state_to_node[state] = node + + if not disable_backoff and not end: + # Add back off when adding node. + for n in range(1, len(state) + 1): + backoff_node = state_to_node.get(state[n:], None) + + # Epsilon transition to the back-off state. + if backoff_node is not None: + graph.add_arc(node, backoff_node, gtn.epsilon) + break + return node + + for grams in ngrams: + for gram in grams: + istate, ostate = gram[:-1], gram[len(gram) - ngram + 1 :] + inode = get_node(istate) + + if END_IDX not in gram[1:] and gram[1:] not in state_to_node: + raise ValueError( + "Ill formed counts: if (x, y_1, ..., y_{n-1}) is above" + "the n-gram threshold, then (y_1, ..., y_{n-1}) must be" + "above the (n-1)-gram threshold" + ) + + if END_IDX in ostate: + # Merge all state having </s> into one as final graph generated + # will be similar. + ostate = tuple([END_IDX]) + + onode = get_node(ostate) + # p(gram[-1] | gram[:-1]) + graph.add_arc( + inode, onode, gtn.epsilon if gram[-1] == END_IDX else gram[-1] + ) + return graph + + +def count_ngrams(lines: List, ngram: List, tokens_to_index: Dict) -> List: + """Counts the number of ngrams.""" + counts = [collections.Counter() for _ in range(ngram)] + for line in lines: + # Prepend implicit start token. + token_line = [START_IDX] + for t in line: + token_line.append(tokens_to_index[t]) + token_line.append(END_IDX) + for n, counter in enumerate(counts): + start_offset = n == 0 + end_offset = ngram == 1 + for e in range(n + start_offset, len(token_line) - end_offset): + counter[tuple(token_line[e - n : e + 1])] += 1 + + return counts + + +def prune_ngrams(ngrams: List, prune: List) -> List: + """Prunes ngrams.""" + pruned_ngrams = [] + for n, grams in enumerate(ngrams): + grams = grams.most_common() + pruned_grams = [gram for gram, c in grams if c > prune[n]] + pruned_ngrams.append(pruned_grams) + return pruned_ngrams + + +def add_blank_grams(pruned_ngrams: List, num_tokens: int, blank: str) -> List: + """Adds blank token to grams.""" + all_grams = [gram for grams in pruned_ngrams for gram in grams] + maxorder = len(pruned_ngrams) + blank_grams = {} + if blank == "forced": + pruned_ngrams = [pruned_ngrams[0] if i == 0 else [] for i in range(maxorder)] + pruned_ngrams[0].append(tuple([num_tokens])) + blank_grams[tuple([num_tokens])] = True + + for gram in all_grams: + # Iterate over all possibilities by using a vector of 0s, 1s to + # denote whether a blank is being used at each position. + if blank == "optional": + # Given a gram ab.. if order n, we have n + 1 positions + # available whether to use blank or not. + onehot_vectors = itertools.product([0, 1], repeat=len(gram) + 1) + elif blank == "forced": + # Must include a blank token in between. + onehot_vectors = [[1] * (len(gram) + 1)] + else: + raise ValueError( + "Invalid value specificed for blank. Must be in |optional|forced|none|" + ) + + for j in onehot_vectors: + new_array = [] + for idx, oz in enumerate(j[:-1]): + if oz == 1 and gram[idx] != START_IDX: + new_array.append(num_tokens) + new_array.append(gram[idx]) + if j[-1] == 1 and gram[-1] != END_IDX: + new_array.append(num_tokens) + for n in range(maxorder): + for e in range(n, len(new_array)): + cur_gram = tuple(new_array[e - n : e + 1]) + if num_tokens in cur_gram and cur_gram not in blank_grams: + pruned_ngrams[n].append(cur_gram) + blank_grams[cur_gram] = True + + return pruned_ngrams + + +def add_self_loops(pruned_ngrams: List) -> List: + """Adds self loops to the ngrams.""" + maxorder = len(pruned_ngrams) + + # Use dict for fast search. + all_grams = set([gram for grams in pruned_ngrams for gram in grams]) + for o in range(1, maxorder): + for gram in pruned_ngrams[o - 1]: + # Repeat one of the tokens. + for pos in range(len(gram)): + if gram[pos] == START_IDX or gram[pos] == END_IDX: + continue + new_gram = gram[:pos] + (gram[pos],) + gram[pos:] + + if new_gram not in all_grams: + pruned_ngrams[o].append(new_gram) + all_grams.add(new_gram) + return pruned_ngrams + + +def parse_lines(lines: List, lexicon: Path) -> List: + """Parses lines with a lexicon.""" + with open(lexicon, "r") as f: + lex = (line.strip().split() for line in f) + lex = {line[0]: line[1:] for line in lex} + print(len(lex)) + return [[t for w in line.split(WORDSEP) for t in lex[w]] for line in lines] + + +@click.command() +@click.option("--data_dir", type=str, default=None, help="Path to dataset root.") +@click.option( + "--tokens", type=str, help="Path to token list (in order used with training)." +) +@click.option("--lexicon", type=str, default=None, help="Path to lexicon") +@click.option( + "--prune", + nargs=2, + type=int, + help="Threshold values for prune unigrams, bigrams, etc.", +) +@click.option( + "--blank", + default=click.Choice(["none", "optional", "forced"]), + help="Specifies the usage of blank token" + "'none' - do not use blank token " + "'optional' - allow an optional blank inbetween tokens" + "'forced' - force a blank inbetween tokens (also referred to as garbage token)", +) +@click.option("--self_loops", is_flag=True, help="Add self loops for tokens") +@click.option("--disable_backoff", is_flag=True, help="Disable backoff transitions") +@click.option("--save_path", default=None, help="Path to save transition graph.") +def cli( + data_dir: str, + tokens: str, + lexicon: str, + prune: List[int], + blank: str, + self_loops: bool, + disable_backoff: bool, + save_path: str, +) -> None: + """CLI for creating the transitions.""" + logger.info(f"Building {len(prune)}-gram transition models.") + + if data_dir is None: + data_dir = ( + Path(__file__).resolve().parents[2] / "data" / "processed" / "iam_lines" + ) + logger.debug(f"Using data dir: {data_dir}") + if not data_dir.exists(): + raise RuntimeError(f"Could not locate iamdb directory at {data_dir}") + else: + data_dir = Path(data_dir) + + # Build table of counts and the back-off if below threshold. + with open(data_dir / "train.txt", "r") as f: + lines = [line.strip() for line in f] + + with open(data_dir / tokens, "r") as f: + tokens = [line.strip() for line in f] + + if lexicon is not None: + lexicon = data_dir / lexicon + lines = parse_lines(lines, lexicon) + + tokens_to_idx = {t: e for e, t in enumerate(tokens)} + + ngram = len(prune) + + logger.info("Counting data...") + ngrams = count_ngrams(lines, ngram, tokens_to_idx) + + pruned_ngrams = prune_ngrams(ngrams, prune) + + for n in range(ngram): + logger.info(f"Kept {len(pruned_ngrams[n])} of {len(ngrams[n])} {n + 1}-grams") + + if blank == "none": + pruned_ngrams = add_blank_grams(pruned_ngrams, len(tokens_to_idx), blank) + + if self_loops: + pruned_ngrams = add_self_loops(pruned_ngrams) + + logger.info("Building graph from pruned ngrams...") + graph = build_graph(pruned_ngrams, disable_backoff) + logger.info(f"Graph has {graph.num_arcs()} arcs and {graph.num_nodes()} nodes.") + + save_path = str(data_dir / save_path) + + logger.info(f"Saving graph to {save_path}") + gtn.save(save_path, graph) + + +if __name__ == "__main__": + cli() diff --git a/text_recognizer/data/iam_preprocessor.py b/text_recognizer/data/iam_preprocessor.py index 5d0fad6..a47aeed 100644 --- a/text_recognizer/data/iam_preprocessor.py +++ b/text_recognizer/data/iam_preprocessor.py @@ -166,7 +166,7 @@ def cli( """CLI for extracting text data from the iam dataset.""" if data_dir is None: data_dir = ( - Path(__file__).resolve().parents[3] / "data" / "raw" / "iam" / "iamdb" + Path(__file__).resolve().parents[2] / "data" / "downloaded" / "iam" / "iamdb" ) logger.debug(f"Using data dir: {data_dir}") if not data_dir.exists(): diff --git a/text_recognizer/data/make_wordpieces.py b/text_recognizer/data/make_wordpieces.py new file mode 100644 index 0000000..e062c4c --- /dev/null +++ b/text_recognizer/data/make_wordpieces.py @@ -0,0 +1,114 @@ +"""Creates word pieces from a text file. + +Most code stolen from: + + https://github.com/facebookresearch/gtn_applications/blob/master/scripts/make_wordpieces.py + +""" +import io +from pathlib import Path +from typing import List, Optional, Union + +import click +from loguru import logger +import sentencepiece as spm + +from text_recognizer.data.iam_preprocessor import load_metadata + + +def iamdb_pieces( + data_dir: Path, text_file: str, num_pieces: int, output_prefix: str +) -> None: + """Creates word pieces from the iamdb train text.""" + # Load training text. + with open(data_dir / text_file, "r") as f: + text = [line.strip() for line in f] + + sp = train_spm_model( + iter(text), + num_pieces + 1, # To account for <unk> + user_symbols=["/"], # added so token is in the output set + ) + + vocab = sorted(set(w for t in text for w in t.split("▁") if w)) + if "move" not in vocab: + raise RuntimeError("`MOVE` not in vocab") + + save_pieces(sp, num_pieces, data_dir, output_prefix, vocab) + + +def train_spm_model( + sentences: iter, vocab_size: int, user_symbols: Union[str, List[str]] = "" +) -> spm.SentencePieceProcessor: + """Trains the sentence piece model.""" + model = io.BytesIO() + spm.SentencePieceTrainer.train( + sentence_iterator=sentences, + model_writer=model, + vocab_size=vocab_size, + bos_id=-1, + eos_id=-1, + character_coverage=1.0, + user_defined_symbols=user_symbols, + ) + sp = spm.SentencePieceProcessor(model_proto=model.getvalue()) + return sp + + +def save_pieces( + sp: spm.SentencePieceProcessor, + num_pieces: int, + data_dir: Path, + output_prefix: str, + vocab: set, +) -> None: + """Saves word pieces to disk.""" + logger.info(f"Generating word piece list of size {num_pieces}.") + pieces = [sp.id_to_piece(i) for i in range(1, num_pieces + 1)] + logger.info(f"Encoding vocabulary of size {len(vocab)}.") + encoded_vocab = [sp.encode_as_pieces(v) for v in vocab] + + # Save pieces to file. + with open(data_dir / f"{output_prefix}_tokens_{num_pieces}.txt", "w") as f: + f.write("\n".join(pieces)) + + # Save lexicon to a file. + with open(data_dir / f"{output_prefix}_lex_{num_pieces}.txt", "w") as f: + for v, p in zip(vocab, encoded_vocab): + f.write(f"{v} {' '.join(p)}\n") + + +@click.command() +@click.option("--data_dir", type=str, default=None, help="Path to processed iam dir.") +@click.option( + "--text_file", type=str, default=None, help="Name of sentence piece training text." +) +@click.option( + "--output_prefix", + type=str, + default="word_pieces", + help="Prefix name to store tokens and lexicon.", +) +@click.option("--num_pieces", type=int, default=1000, help="Number of word pieces.") +def cli( + data_dir: Optional[str], + text_file: Optional[str], + output_prefix: Optional[str], + num_pieces: Optional[int], +) -> None: + """CLI for training the sentence piece model.""" + if data_dir is None: + data_dir = ( + Path(__file__).resolve().parents[2] / "data" / "processed" / "iam_lines" + ) + logger.debug(f"Using data dir: {data_dir}") + if not data_dir.exists(): + raise RuntimeError(f"Could not locate iamdb directory at {data_dir}") + else: + data_dir = Path(data_dir) + + iamdb_pieces(data_dir, text_file, num_pieces, output_prefix) + + +if __name__ == "__main__": + cli() |