# 導(dǎo)入第三方的庫函數(shù)
import os
from io import open
import time

import torch
import torch.nn as nn
import torch.quantization
import torch.nn.functional as F

# 創(chuàng)建LSTM模型類
class LSTMModel(nn.Module):
    """整個網(wǎng)絡(luò)包含一個encoder, 一個recurrent模塊和一個decoder."""

    def __init__(self, ntoken, ninp, nhid, nlayers, dropout=0.5):
        super(LSTMModel, self).__init__()
        # 預(yù)定義一些網(wǎng)絡(luò)層
        self.drop = nn.Dropout(dropout)
        # 嵌入層
        self.encoder = nn.Embedding(ntoken, ninp)
        # LSTM層
        self.rnn = nn.LSTM(ninp, nhid, nlayers, dropout=dropout)
        # 線性層
        self.decoder = nn.Linear(nhid, ntoken)
        self.init_weights()
        self.nhid = nhid
        self.nlayers = nlayers

    def init_weights(self):
     '''
     初始化模型權(quán)重
     '''
        initrange = 0.1
        self.encoder.weight.data.uniform_(-initrange, initrange)
        self.decoder.bias.data.zero_()
        self.decoder.weight.data.uniform_(-initrange, initrange)

    def forward(self, input, hidden):
     '''
     搭建網(wǎng)絡(luò)并執(zhí)行前向推理
     '''
        emb = self.drop(self.encoder(input))
        output, hidden = self.rnn(emb, hidden)
        output = self.drop(output)
        decoded = self.decoder(output)
        return decoded, hidden

    def init_hidden(self, bsz):
     '''
     初始化hidden層的權(quán)重
     '''
        weight = next(self.parameters())
        return (weight.new_zeros(self.nlayers, bsz, self.nhid),
                weight.new_zeros(self.nlayers, bsz, self.nhid))

# 創(chuàng)建一個詞典類，用來處理數(shù)據(jù)
# 構(gòu)建詞匯表，包括詞到索引的映射和索引到詞的映射
class Dictionary(object):
    def __init__(self):
        self.word2idx = {}
        self.idx2word = []

    def add_word(self, word):
     '''
     在詞典中添加新的word
     '''
        if word not in self.word2idx:
            self.idx2word.append(word)
            self.word2idx[word] = len(self.idx2word) - 1
        return self.word2idx[word]

    def __len__(self):
     '''
     返回詞典的長度
     '''
        return len(self.idx2word)

#  Corpus 類:處理文本數(shù)據(jù)，包括讀取文件、構(gòu)建詞匯表和將文本轉(zhuǎn)換為索引序列
class Corpus(object):
    def __init__(self, path):
        self.dictionary = Dictionary()
        # 分別獲取訓(xùn)練集、驗(yàn)證集和測試集
        self.train = self.tokenize(os.path.join(path, 'train.txt'))
        self.valid = self.tokenize(os.path.join(path, 'valid.txt'))
        self.test = self.tokenize(os.path.join(path, 'test.txt'))

    def tokenize(self, path):
        """對輸入的文件執(zhí)行分詞操作"""
        assert os.path.exists(path)
        # 將新的單詞添加到詞典中
        with open(path, 'r', encoding="utf8") as f:
            for line in f:
                words = line.split() + ['<eos>']
                for word in words:
                    self.dictionary.add_word(word)

        # 標(biāo)記文件的內(nèi)容
        with open(path, 'r', encoding="utf8") as f:
            idss = []
            for line in f:
                words = line.split() + ['<eos>']
                ids = []
                for word in words:
                    ids.append(self.dictionary.word2idx[word])
                idss.append(torch.tensor(ids).type(torch.int64))
            ids = torch.cat(idss)
        return ids

# 設(shè)置模型的路徑
model_data_filepath = 'data/'
corpus = Corpus(model_data_filepath + 'wikitext-2')
ntokens = len(corpus.dictionary)

# 搭建網(wǎng)絡(luò)模型
model = LSTMModel(
    ntoken = ntokens,
    ninp = 512,
    nhid = 256,
    nlayers = 5,
)

# 加載預(yù)訓(xùn)練的模型權(quán)重
model.load_state_dict(
    torch.load(
        model_data_filepath + 'word_language_model_quantize.pth',
        map_location=torch.device('cpu')
        )
    )
# 將模型切換為推理模式，并打印整個模型
model.eval()
print(model)

# 獲取一個隨機(jī)的輸入數(shù)值
input_ = torch.randint(ntokens, (1, 1), dtype=torch.long)
hidden = model.init_hidden(1)
temperature = 1.0
num_words = 1000

# 遍歷數(shù)據(jù)集進(jìn)行前向推理并將結(jié)果保存起來
with open(model_data_filepath + 'out.txt', 'w') as outf:
    with torch.no_grad():  # no tracking history
        for i in range(num_words):
            output, hidden = model(input_, hidden)
            word_weights = output.squeeze().div(temperature).exp().cpu()
            word_idx = torch.multinomial(word_weights, 1)[0]
            input_.fill_(word_idx)
            word = corpus.dictionary.idx2word[word_idx]
            outf.write(str(word.encode('utf-8')) + ('\n' if i % 20 == 19 else ' '))
            if i % 100 == 0:
                print('| Generated {}/{} words'.format(i, 1000))

with open(model_data_filepath + 'out.txt', 'r') as outf:
    all_output = outf.read()
    print(all_output)

bptt = 25
criterion = nn.CrossEntropyLoss()
eval_batch_size = 1

# 創(chuàng)建測試數(shù)據(jù)集
def batchify(data, bsz):
    # 對測試數(shù)據(jù)集進(jìn)行分塊
    nbatch = data.size(0) // bsz
    # 去掉多余的元素
    data = data.narrow(0, 0, nbatch * bsz)
    # 在bsz批處理中平均劃分?jǐn)?shù)據(jù)
    return data.view(bsz, -1).t().contiguous()

test_data = batchify(corpus.test, eval_batch_size)

# 獲取bath塊的輸入數(shù)據(jù)
def get_batch(source, i):
    seq_len = min(bptt, len(source) - 1 - i)
    data = source[i:i+seq_len]
    target = source[i+1:i+1+seq_len].view(-1)
    return data, target

def repackage_hidden(h):
  """
  用新的張量把隱藏的狀態(tài)包裝起來，把它們從歷史中分離出來
  """
  if isinstance(h, torch.Tensor):
      return h.detach()
  else:
      return tuple(repackage_hidden(v) for v in h)
# 評估函數(shù)
def evaluate(model_, data_source):
    # 打開評估模式
    model_.eval()
    total_loss = 0.
    hidden = model_.init_hidden(eval_batch_size)
    with torch.no_grad():
        for i in range(0, data_source.size(0) - 1, bptt):
         # 獲取測試數(shù)據(jù)
            data, targets = get_batch(data_source, i)
            # 執(zhí)行前向推理
            output, hidden = model_(data, hidden)
            hidden = repackage_hidden(hidden)
            output_flat = output.view(-1, ntokens)
            # 獲取訓(xùn)練loss
            total_loss += len(data) * criterion(output_flat, targets).item()
    return total_loss / (len(data_source) - 1)

# 初始化動態(tài)量化模塊
quantized_model = torch.quantization.quantize_dynamic(
    model, {nn.LSTM, nn.Linear}, dtype=torch.qint8
)
print(quantized_model)

def print_size_of_model(model):
    torch.save(model.state_dict(), "temp.p")
    print('Size (MB):', os.path.getsize("temp.p")/1e6)
    os.remove('temp.p')

print_size_of_model(model)
print_size_of_model(quantized_model)

torch.set_num_threads(1)
# 評估模型的運(yùn)行時間
def time_model_evaluation(model, test_data):
    s = time.time()
    loss = evaluate(model, test_data)
    elapsed = time.time() - s
    print('''loss: {0:.3f}\nelapsed time (seconds): {1:.1f}'''.format(loss, elapsed))

time_model_evaluation(model, test_data)
time_model_evaluation(quantized_model, test_data)