BigramLanguageModel/main.py at main · Joelcic/BigramLanguageModel · 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
from utils import load_txt, load_json
from tokenization import *
from model import BigramLanguageModel
from data_preperation import add_BOS_EOS

if __name__ == "__main__":

    #Load or process the text
    """# Data Preperation
    text_path = "data/frankenstein.txt"
    processed_txt = add_BOS_EOS(file_path=text_path)"""

    # Load already processed text
    text_path_processed = "artifacts/frankenstein_processed.txt"
    processed_txt = load_txt(filepath=text_path_processed)


    # Tokenize and create mapping
    create_word_tok_map(text_path_processed)
    stoi_path = "artifacts/string_to_int_word.json"
    itos_path = "artifacts/int_to_string_word.json"

    # Boundry tokens
    bos_char =  "§"
    bos_id = word_encode(bos_char, stoi_path)
    eos_char = "¶"
    eos_id = word_encode(eos_char, stoi_path)

    # Create the model
    vocab_size = len(load_json(stoi_path))
    bigramLM = BigramLanguageModel(vocab_size)

    # Calculate teh distribution
    text_encoded = word_encode(processed_txt, stoi_path)
    bigramLM.fit(text_encoded=text_encoded, alpha=1/vocab_size)

    # Generate text
    generated_text = bigramLM.sample(n=200, bos_id=bos_id[0], eos_id=eos_id[0])
    generated_text_decoded = word_decode(generated_text, itos_path, sep=" ")

    print(generated_text_decoded)