# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: 'Python 3.9.2 64-bit (''ocr-uva'': conda)'
#     name: python392jvsc74a57bd0cc29f658ddb1b0f0a648f4c47acf5938bc6d1ad3f68ae93354e191176a755a49
# ---

# +
from gensim.models import Word2Vec

import os
import re
import string

from nltk.tokenize import RegexpTokenizer


# +
# Load data

newseye_path = os.path.join('..', 'data', 'newseye')
icdar_2017_1_path = os.path.join(newseye_path, '2017', 'full', 'eng_monograph')
icdar_2017_2_path = os.path.join(newseye_path, '2017', 'full', 'eng_periodical')
icdar_2019_path = os.path.join(newseye_path, '2019', 'full', 'EN')

# -

tokenizer = RegexpTokenizer(r'\w+')

# +
documents = []

for icdar_path in [icdar_2017_1_path, icdar_2017_2_path, icdar_2019_path]:
    for filename in os.listdir(icdar_path):
        file_path = os.path.join(icdar_path, filename)
        with open(file_path, 'r', encoding='utf-8') as text_file:
            file_lines = text_file.readlines()
            gt_line = file_lines[2]
            processed_line = gt_line.replace('[ GS_aligned]', '').replace('#', '').replace('@', '')

            text_nonum = re.sub(r'\d+', '', processed_line)
            text_nopunct = "".join([char.lower() for char in text_nonum if char not in string.punctuation])
            text_no_doublespace = re.sub('\s+', ' ', text_nopunct).strip()
            result = tokenizer.tokenize(text_no_doublespace)
            documents.append(result)
# -

model_path = 'gensim_default_eng.model'


def load_model():
    if not os.path.exists(model_path):
        return None

    model = Word2Vec.load(model_path)
    return model



# +
# TRAIN

def create_model(corpus):
    model = Word2Vec(vector_size=300, window=5, min_count=5, workers=2)
    model.build_vocab(corpus, progress_per=10000)
    model.train(corpus, total_examples=model.corpus_count, epochs=300, report_delay=1)
    model.save()
    return model


# -

model = load_model()
if model is None:
    print('Model is not loaded. Creating and training now...')
    model = create_model(documents)

words = ['man', 'new', 'time', 'day', 'good', 'old', 'little', 'one', 'two', 'three']
for word in words:
    print(f'-- \'{word}\':')
    print(model.wv.most_similar(positive=[word]))



StarcoderdataJupyterScriptsDedupFiltered