# Path: data/common.py
# WORD_EMB_PATH = '/path/to/word/embedding/'
# 
# Path: data/bookcorpus.py
# VOCAB_PATH = os.path.join(BOOKCORPUS_DATA_DIR, 'word_count{}')


from gensim.models import KeyedVectors
from data.common import WORD_EMB_PATH
from data.bookcorpus import VOCAB_PATH
import collections
import os
import tqdm
import numpy as np
import sklearn.linear_model
import tensorflow as tf
  # Train a linear regression model on the shared embedding vectors.
  print("Training linear regression model")
  model = sklearn.linear_model.LinearRegression()
  model.fit(shared_w2v_emb, shared_st_emb)

  # Create the expanded vocabulary.
  print("Creating embeddings for expanded vocabuary")
  combined_emb = collections.OrderedDict()
  # Ignore words with underscores (spaces).
  all_words = [w for w in word2vec.vocab if "_" not in w]
  all_words = np.asarray(all_words)
  num_words = len(all_words)
  batch_size = 2048
  num_batches = num_words // batch_size + (num_words % batch_size != 0)

  for batch_idx in tqdm.trange(num_batches):
    ws = all_words[batch_idx * batch_size: (batch_idx + 1) * batch_size]
    w_embs = model.predict(word2vec[ws])
    for w, w_emb in zip(ws, w_embs):
      combined_emb[w] = w_emb

  for w in vocab:
    combined_emb[w] = thought_emb[vocab[w]]

  print("Created expanded vocabulary of %d words" % len(combined_emb))
  return combined_emb


def load_bookcorpus_vocab():
  vocab = dict()

# Path: data/common.py
# WORD_EMB_PATH = '/path/to/word/embedding/'
# 
# Path: data/bookcorpus.py
# VOCAB_PATH = os.path.join(BOOKCORPUS_DATA_DIR, 'word_count{}')


# Path: thought/vocabulary_expansion.py
from gensim.models import KeyedVectors
from data.common import WORD_EMB_PATH
from data.bookcorpus import VOCAB_PATH
import collections
import os
import tqdm
import numpy as np
import sklearn.linear_model
import tensorflow as tf

  # Train a linear regression model on the shared embedding vectors.
  print("Training linear regression model")
  model = sklearn.linear_model.LinearRegression()
  model.fit(shared_w2v_emb, shared_st_emb)

  # Create the expanded vocabulary.
  print("Creating embeddings for expanded vocabuary")
  combined_emb = collections.OrderedDict()
  # Ignore words with underscores (spaces).
  all_words = [w for w in word2vec.vocab if "_" not in w]
  all_words = np.asarray(all_words)
  num_words = len(all_words)
  batch_size = 2048
  num_batches = num_words // batch_size + (num_words % batch_size != 0)

  for batch_idx in tqdm.trange(num_batches):
    ws = all_words[batch_idx * batch_size: (batch_idx + 1) * batch_size]
    w_embs = model.predict(word2vec[ws])
    for w, w_emb in zip(ws, w_embs):
      combined_emb[w] = w_emb

  for w in vocab:
    combined_emb[w] = thought_emb[vocab[w]]

  print("Created expanded vocabulary of %d words" % len(combined_emb))
  return combined_emb


def load_bookcorpus_vocab():
  vocab = dict()

  with open(VOCAB_PATH) as f: