OCR_Correction/nlc_data.py at master · hoangtuanvu/OCR_Correction · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
# -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import re
import os
import sys
import codecs
from collections import defaultdict
from tensorflow.python.platform import gfile

# Special vocabulary symbols - we always put them at the start.
_PAD = "<pad>"
_SOS = "<sos>"
_EOS = "<eos>"
_UNK = "<unk>"
_START_VOCAB = [_PAD, _SOS, _EOS, _UNK]

PAD_ID = 0
SOS_ID = 1
EOS_ID = 2
UNK_ID = 3

# Regular expressions used to tokenize.
_DIGIT_RE = re.compile(r"\d")


# char tokenizer: split the sentence into a list of tokens.
# sentence is `unicode object`
def char_tokenizer(sentence):
    return list(sentence.strip())


# create vocabulary from train&dev data
def create_vocabulary(vocabulary_path, data_paths, max_vocabulary_size, tokenizer=None, normalize_digits=False):
    if not os.path.exists(vocabulary_path):
        print("Creating vocabulary %s from data %s" % (vocabulary_path, str(data_paths)))
        vocab = defaultdict(int)
        for path in data_paths:
            with codecs.open(path, mode="r", encoding="utf-8") as fr:
                counter = 0
                for line in fr:
                    counter += 1
                    if counter % 100000 == 0:
                        print("  processing line %d" % (counter,))
                    tokens = tokenizer(line)
                    for w in tokens:
                        word = re.sub(_DIGIT_RE, "0", w) if normalize_digits else w
                        vocab[word] += 1

        vocab_list = _START_VOCAB + sorted(vocab, key=vocab.get, reverse=True)
        print("Vocabulary size: %d" % len(vocab_list))
        if len(vocab_list) > max_vocabulary_size:
            vocab_list = vocab_list[:max_vocabulary_size]
        with codecs.open(vocabulary_path, mode="w", encoding="utf-8") as vocab_file:
            for w in vocab_list:
                vocab_file.write(w + "\n")


def initialize_vocabulary(vocabulary_path):
    if os.path.exists(vocabulary_path):
        rev_vocab = []
        with codecs.open(vocabulary_path, mode="r", encoding="utf-8") as fr:
            rev_vocab.extend(fr.readlines())
        rev_vocab = [line.strip("\n") for line in rev_vocab]
        vocab = dict([(x, y) for (y, x) in enumerate(rev_vocab)])
        return vocab, rev_vocab
    else:
        raise ValueError("Vocabulary file %s not found.", vocabulary_path)


def sentence_to_token_ids(sentence, vocabulary, tokenizer=None, normalize_digits=False):
    words = tokenizer(sentence)
    if not normalize_digits:
        return [vocabulary.get(w, UNK_ID) for w in words]
    # Normalize digits by 0 before looking words up in the vocabulary.
    return [vocabulary.get(re.sub(_DIGIT_RE, "0", w), UNK_ID) for w in words]


def data_to_token_ids(data_path, target_path, vocabulary_path,
                      tokenizer=None, normalize_digits=False):
    if not os.path.exists(target_path):
        print("Tokenizing data in %s" % (data_path,))
        vocab, _ = initialize_vocabulary(vocabulary_path)
        with codecs.open(data_path, mode="r", encoding="utf-8") as data_file:
            with gfile.GFile(target_path, mode="w") as tokens_file:
                for idx, line in enumerate(data_file):
                    if idx + 1 % 100000 == 0:
                        print("  tokenizing line %d" % idx)
                    token_ids = sentence_to_token_ids(line, vocab, tokenizer, normalize_digits)
                    tokens_file.write(" ".join([str(tok) for tok in token_ids]) + "\n")


def get_data_path(data_dir, data_type="train"):
    data_path = os.path.join(data_dir, data_type)
    print("[%s] x data path is:%s" % (data_type.upper(), data_path + ".x.txt"))
    print("[%s] y data path is:%s" % (data_type.upper(), data_path + ".y.txt"))
    if os.path.exists(data_path + ".x.txt") and os.path.exists(data_path + ".y.txt"):
        return data_path
    else:
        print("Please provider [%s] x&y data" % (data_type.upper()))
        sys.exit()


def prepare_nlc_data(data_dir, max_vocabulary_size, tokenizer=char_tokenizer):
    # Get nlc data to the specified directory.
    train_path = get_data_path(data_dir, "train")
    dev_path = get_data_path(data_dir, "valid")

    # Create vocabularies of the appropriate sizes.
    vocab_path = os.path.join(data_dir, "vocab.dat")
    create_vocabulary(vocab_path, [train_path + ".y.txt", train_path + ".x.txt"], max_vocabulary_size, tokenizer)

    # Create token ids for the training data.
    y_train_ids_path = train_path + ".ids.y"
    x_train_ids_path = train_path + ".ids.x"
    data_to_token_ids(train_path + ".y.txt", y_train_ids_path, vocab_path, tokenizer)
    data_to_token_ids(train_path + ".x.txt", x_train_ids_path, vocab_path, tokenizer)

    # Create token ids for the development data.
    y_dev_ids_path = dev_path + ".ids.y"
    x_dev_ids_path = dev_path + ".ids.x"
    data_to_token_ids(dev_path + ".y.txt", y_dev_ids_path, vocab_path, tokenizer)
    data_to_token_ids(dev_path + ".x.txt", x_dev_ids_path, vocab_path, tokenizer)

    return x_train_ids_path, y_train_ids_path, x_dev_ids_path, y_dev_ids_path, vocab_path