LanguageModel/format_data.py at main · yuvaltimen/LanguageModel · 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
import sys
from nltk import sent_tokenize
from random import shuffle

def main():
    """
    Tokenizes the text by sentence, removes undesirable punctuation, and saves
    the file with correct number of SENT_BEGIN and SENT_END tokens.
    Parameters:
        None
    Returns:
        None
    """
    with open(sys.argv[1], "r") as to_format:
        text = to_format.read()

    sentences = sent_tokenize(text)
    for idx in range(len(sentences)):
        sentences[idx] = clean_sentence(sentences[idx])

    shuffle(sentences)

    with open(sys.argv[2], "w+") as to_write:
        prepend = ["<s>" for _ in range(int(sys.argv[3]) - 1)]
        postpend = ["</s>" for _ in range(int(sys.argv[3]) - 1)]
        for sent in sentences:
            to_write.write(" ".join(prepend) + " " + sent + " " + " ".join(postpend) + "\n")


def clean_sentence(sent):
    """Removes undesirable punctuation from the given sentence.
    Parameters:
        sent (str) The string sentence to clean.
    Returns:
        str: The cleaned sentence.
    """
    sent = sent.replace(".", "")
    sent = sent.replace(",", "")
    sent = sent.replace(":", "")
    sent = sent.replace(";", "")
    sent = sent.replace("!", "")
    sent = sent.replace("?", "")
    sent = sent.replace("<", "")
    sent = sent.replace(">", "")
    sent = sent.replace("-", "")
    sent = sent.replace("/", "")
    sent = sent.replace("\\", "")
    sent = sent.replace("{", "")
    sent = sent.replace("}", "")
    sent = sent.replace("%", "")
    sent = sent.replace("|", "")
    sent = sent.replace("_", "")
    sent = sent.replace(".", "")
    sent = sent.replace("”", "")
    sent = sent.replace("“", "")
    sent = sent.replace("\n", "")
    return sent


if __name__ == '__main__':

    if len(sys.argv) != 4:
        raise Exception("Usage: python3 format_data.py [unformatted_file_path] [file_save_path] [n-gram_order]")
    main()