# coding: utf-8
from dataset import ptb
from rnnlm_gen import RnnlmGen
import sys
sys.path.append('..')


# データ読み込み
corpus, word_to_id, id_to_word = ptb.load_data('train')
# 語彙数（単語の種類数）を取得 => 10000
vocab_size = len(word_to_id)
# 単語数を取得 => 929589
corpus_size = len(corpus)

# モデル定義
model = RnnlmGen()
model.load_params('../ch06/Rnnlm.pkl')

# 最初の単語
start_word = 'you'  # 先頭の単語をyouとする
# 最初の単語の単語IDを取得
start_id = word_to_id[start_word]
# skip文字の設定
skip_words = ['N', '<unk>', '$']
skip_ids = [word_to_id[w] for w in skip_words]

# 文章生成（最初の一単語を与えて，複数単語を得る）
# e.g. word_ids = [316, 3371, 2053, 7103, 809, 8342, 4600, 9117, 2670, 3503]
word_ids = model.generate(start_id, skip_ids)

# 取得したword_idsをidと単語のmappingを利用して単語に変換・結合する
# e.g. ' '.join(['you', 'say', 'goodbye'])
txt = ' '.join([id_to_word[i] for i in word_ids])
# 文章区切りを表す<eos>を，ピリオドと改行に変換
txt = txt.replace(' <eos>', '.\n')

print(txt)

# 出力例：
# you taxpayers notion king watches demonstrated alexander comex credibility consist
