w2v_vocab_randomnly.py

import nltk
from nltk.corpus import stopwords
from gensim.models import Word2Vec,KeyedVectors
from gensim.test.utils import datapath
import re
import glob
from tqdm import tqdm
import gensim
import multiprocessing
import random

#Clean the data function
stopwords_list=stopwords.words('english')
cores=multiprocessing.cpu_count()

def clean_data(w):
    w = w.lower()                                                                             #Lower casing
    w=re.sub(r'[^\w\s]','',w)                                                                 #Remove other special characters
    w=re.sub(r"([0-9])", r" ",w)                                                              #Remove numbers       
    words = w.split() 
    clean_words = [word for word in words if (word not in stopwords_list) and len(word) > 2]  #Remove stop and short words
    return " ".join(clean_words)  

#Load the input function
class SentGen(object):
    def __init__(self, dirname,input_mode):                   #dirname is the path to multiple .txt files
        self.dirname = dirname
        self.input_mode=input_mode
    def __iter__(self):
        for idx,fname in enumerate(self.dirname,):
            for line in get_sentences(fname,self.input_mode):       #Get one line from get_sentences and yield the tokens by line.split()
                yield line.split()

def get_sentences(fname,input_mode):
    with open(fname,'r') as f:
            texts=f.readlines()
    if input_mode=='vocab':
      texts=random.sample(texts,texts_sample)      # Randomnly selecting a sample of sentences from a file if input_mode is 'vocab'
    sent=list(map(clean_data,texts))               # Clean the sentences
    for line in tqdm(sent):                        # Read one line and yield
      return lines
    
#Give input
input_data=glob.glob('path_to_folder_with_multiple_txt_files_inside_it')

#Initialize the model
model = Word2Vec(min_count=1,window=3,size=300,workers=cores-1)

#Build vocabulary
vocab_files=random.sample(input_data,vocab_sample)                  # Randomnly selecting a sample of files from all the files
vocab_inps=SentGen(vocab_files,'vocab')                             # 'vocab' is to take random sentences from the files.
model.build_vocab(vocab_inps)                                       # Building the vocabulary using entire dataset
vocab_len=len(model.wv.vocab)

#Training the model
model_input=SentGen(input_data,'all_data')
model.train(model_input,total_examples=model.corpus_count,epochs=50)  #Add callbacks, if required

#Saving the model
model.wv.save_word2vec_format('gensim_w2v_model.bin',binary=True)

#Loading the model during testing time
trained_model= KeyedVectors.load_word2vec_format('gensim_w2v_model.bin', binary=True)