first version

.gitignore

@@ -0,0 +1,22 @@
+# PyTorch
+/.torch
+
+# Data files
+*.csv
+*.json
+*.tsv
+
+# Model files
+*.ckpt
+*.pth
+*.pkl
+
+# Logs and checkpoints
+logs/
+checkpoints/
+
+# Secondary files
+*.pyc
+__pycache__/
+.DS_Store
+

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/misc.xml

@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.9" project-jdk-type="Python SDK" />
+</project>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/nlp.iml" filepath="$PROJECT_DIR$/.idea/nlp.iml" />
+    </modules>
+  </component>
+</project>

.idea/nlp.iml

@@ -0,0 +1,12 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+  <component name="PyDocumentationSettings">
+    <option name="format" value="PLAIN" />
+    <option name="myDocStringFormat" value="Plain" />
+  </component>
+</module>

.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="$PROJECT_DIR$" vcs="Git" />
+  </component>
+</project>

.idea/workspace.xml

@@ -0,0 +1,111 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ChangeListManager">
+    <list default="true" id="975e88fb-d387-4d2c-9625-dc69f610d124" name="Changes" comment="">
+      <change afterPath="$PROJECT_DIR$/.idea/inspectionProfiles/profiles_settings.xml" afterDir="false" />
+      <change afterPath="$PROJECT_DIR$/.idea/misc.xml" afterDir="false" />
+      <change afterPath="$PROJECT_DIR$/.idea/modules.xml" afterDir="false" />
+      <change afterPath="$PROJECT_DIR$/.idea/nlp.iml" afterDir="false" />
+      <change afterPath="$PROJECT_DIR$/.idea/vcs.xml" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/src/apis/train.py" beforeDir="false" afterPath="$PROJECT_DIR$/src/apis/train.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/src/datasets/dataloader.py" beforeDir="false" afterPath="$PROJECT_DIR$/src/datasets/dataloader.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/src/models/LSTM/model.py" beforeDir="false" afterPath="$PROJECT_DIR$/src/models/LSTM/model.py" afterDir="false" />
+    </list>
+    <option name="SHOW_DIALOG" value="false" />
+    <option name="HIGHLIGHT_CONFLICTS" value="true" />
+    <option name="HIGHLIGHT_NON_ACTIVE_CHANGELIST" value="false" />
+    <option name="LAST_RESOLUTION" value="IGNORE" />
+  </component>
+  <component name="Git.Settings">
+    <option name="RECENT_GIT_ROOT_PATH" value="$PROJECT_DIR$" />
+  </component>
+  <component name="MarkdownSettingsMigration">
+    <option name="stateVersion" value="1" />
+  </component>
+  <component name="ProjectId" id="2Q8D9XoYiTKL5jiaHLTd3rsHf4Y" />
+  <component name="ProjectViewState">
+    <option name="hideEmptyMiddlePackages" value="true" />
+    <option name="showLibraryContents" value="true" />
+  </component>
+  <component name="PropertiesComponent">{
+  &quot;keyToString&quot;: {
+    &quot;RunOnceActivity.OpenProjectViewOnStart&quot;: &quot;true&quot;,
+    &quot;RunOnceActivity.ShowReadmeOnStart&quot;: &quot;true&quot;,
+    &quot;last_opened_file_path&quot;: &quot;D:/YOU/dasanxia/NLP/new0522/nlp&quot;,
+    &quot;settings.editor.selected.configurable&quot;: &quot;com.jetbrains.python.configuration.PyActiveSdkModuleConfigurable&quot;
+  }
+}</component>
+  <component name="RunManager" selected="Python.run_gradio">
+    <configuration name="inference" type="PythonConfigurationType" factoryName="Python" temporary="true" nameIsGenerated="true">
+      <module name="nlp" />
+      <option name="INTERPRETER_OPTIONS" value="" />
+      <option name="PARENT_ENVS" value="true" />
+      <envs>
+        <env name="PYTHONUNBUFFERED" value="1" />
+      </envs>
+      <option name="SDK_HOME" value="" />
+      <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$" />
+      <option name="IS_MODULE_SDK" value="true" />
+      <option name="ADD_CONTENT_ROOTS" value="true" />
+      <option name="ADD_SOURCE_ROOTS" value="true" />
+      <option name="SCRIPT_NAME" value="$PROJECT_DIR$/inference.py" />
+      <option name="PARAMETERS" value="" />
+      <option name="SHOW_COMMAND_LINE" value="false" />
+      <option name="EMULATE_TERMINAL" value="false" />
+      <option name="MODULE_MODE" value="false" />
+      <option name="REDIRECT_INPUT" value="false" />
+      <option name="INPUT_FILE" value="" />
+      <method v="2" />
+    </configuration>
+    <configuration name="run_gradio" type="PythonConfigurationType" factoryName="Python" temporary="true" nameIsGenerated="true">
+      <module name="nlp" />
+      <option name="INTERPRETER_OPTIONS" value="" />
+      <option name="PARENT_ENVS" value="true" />
+      <envs>
+        <env name="PYTHONUNBUFFERED" value="1" />
+      </envs>
+      <option name="SDK_HOME" value="" />
+      <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$" />
+      <option name="IS_MODULE_SDK" value="true" />
+      <option name="ADD_CONTENT_ROOTS" value="true" />
+      <option name="ADD_SOURCE_ROOTS" value="true" />
+      <option name="SCRIPT_NAME" value="$PROJECT_DIR$/run_gradio.py" />
+      <option name="PARAMETERS" value="" />
+      <option name="SHOW_COMMAND_LINE" value="false" />
+      <option name="EMULATE_TERMINAL" value="false" />
+      <option name="MODULE_MODE" value="false" />
+      <option name="REDIRECT_INPUT" value="false" />
+      <option name="INPUT_FILE" value="" />
+      <method v="2" />
+    </configuration>
+    <recent_temporary>
+      <list>
+        <item itemvalue="Python.run_gradio" />
+        <item itemvalue="Python.inference" />
+      </list>
+    </recent_temporary>
+  </component>
+  <component name="SpellCheckerSettings" RuntimeDictionaries="0" Folders="0" CustomDictionaries="0" DefaultDictionary="application-level" UseSingleDictionary="true" transferred="true" />
+  <component name="TaskManager">
+    <task active="true" id="Default" summary="Default task">
+      <changelist id="975e88fb-d387-4d2c-9625-dc69f610d124" name="Changes" comment="" />
+      <created>1684726163448</created>
+      <option name="number" value="Default" />
+      <option name="presentableId" value="Default" />
+      <updated>1684726163448</updated>
+    </task>
+    <servers />
+  </component>
+  <component name="Vcs.Log.Tabs.Properties">
+    <option name="TAB_STATES">
+      <map>
+        <entry key="MAIN">
+          <value>
+            <State />
+          </value>
+        </entry>
+      </map>
+    </option>
+  </component>
+</project>

app.py

@@ -1,7 +1,14 @@
-import gradio as gr
+# !/user/bin/env python3
+# -*- coding: utf-8 -*-
 
-def greet(name):
-    return "Hello " + name + "!!"
+import gradio
+from inference import infer
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+
+
+
+INTERFACE = gradio.Interface(fn=infer, inputs=["text","text"], outputs=["text"], title="Poem Generation",
+                             description="model: lstm/GRU/Seq2Seq/Transformer/GPT-2",
+                             thumbnail="https://github.com/gradio-app/gpt-2/raw/master/screenshots/interface.png?raw=true")
+
+INTERFACE.launch(inbrowser=True)

inference.py

@@ -0,0 +1,112 @@
+import torch
+import argparse
+import numpy as np
+from src.models.LSTM.model import Poetry_Model_lstm
+from src.datasets.dataloader import train_vec
+from src.utils.utils import make_cuda
+from src.models.Transformer.model import Poetry_Model_Transformer
+
+def parse_arguments():
+    # argument parsing
+    parser = argparse.ArgumentParser(description="Specify Params for Experimental Setting")
+    parser.add_argument('--model', type=str, default='lstm',
+                        help="lstm/GRU/Seq2Seq/Transformer/GPT-2")
+    parser.add_argument('--Word2Vec', default=True)
+    parser.add_argument('--strict_dataset', default=False, help="strict dataset")
+    parser.add_argument('--n_hidden', type=int, default=128)
+
+    parser.add_argument('--save_path', type=str, default='save_models/model_params.pth')
+
+    return parser.parse_args()
+
+
+def generate_poetry(model, head_string, w1, word_2_index, index_2_word,args):
+    print("藏头诗生成中...., {}".format(head_string))
+    poem = ""
+    # 以句子的每一个字为开头生成诗句
+    for head in head_string:
+        if head not in word_2_index:
+            print("抱歉，不能生成以{}开头的诗".format(head))
+            return
+
+        sentence = head
+        max_sent_len = 20
+
+        h_0 = torch.tensor(np.zeros((2, 1, args.n_hidden), dtype=np.float32))
+        c_0 = torch.tensor(np.zeros((2, 1, args.n_hidden), dtype=np.float32))
+
+        input_eval = word_2_index[head]
+        for i in range(max_sent_len):
+            if args.Word2Vec:
+                word_embedding = torch.tensor(w1[input_eval][None][None])
+            else:
+                word_embedding = torch.tensor([input_eval]).unsqueeze(dim=0)
+            pre, (h_0, c_0) = model(word_embedding, h_0, c_0)
+            char_generated = index_2_word[int(torch.argmax(pre))]
+
+            if char_generated == '。':
+                break
+            # 以新生成的字为输入继续向下生成
+            input_eval = word_2_index[char_generated]
+            sentence += char_generated
+
+        poem += '\n' + sentence
+
+    return poem
+
+def infer(model,poem_head):
+    args = parse_arguments()
+    args.model=model
+
+    all_data, (w1, word_2_index, index_2_word) = train_vec()
+    args.word_size, args.embedding_num = w1.shape
+    string = poem_head
+    # string = '自然语言'
+
+    if args.model == 'lstm':
+        model = Poetry_Model_lstm(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
+        args.save_path='save_models/lstm_25.pth'
+    elif args.model == 'GRU':
+        model = Poetry_Model_lstm(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
+    elif args.model == 'Seq2Seq':
+        model = Poetry_Model_lstm(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
+    elif args.model == 'Transformer':
+        model = Poetry_Model_Transformer(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
+        args.save_path='save_models/transformer.pth'
+    elif args.model == 'GPT-2':
+        model = Poetry_Model_lstm(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
+    else:
+        print("Please choose a model!\n")
+
+    model.load_state_dict(torch.load(args.save_path))
+    model = make_cuda(model)
+    poem = generate_poetry(model, string, w1, word_2_index, index_2_word,args)
+    return poem
+
+
+if __name__ == '__main__':
+    args = parse_arguments()
+    all_data, (w1, word_2_index, index_2_word) = train_vec()
+    args.word_size, args.embedding_num = w1.shape
+    string = input("诗头:")
+    # string = '自然语言'
+    
+    if args.model == 'lstm':
+        model = Poetry_Model_lstm(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
+        args.save_path='save_models/lstm_25.pth'
+    elif args.model == 'GRU':
+        model = Poetry_Model_lstm(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
+    elif args.model == 'Seq2Seq':
+        model = Poetry_Model_lstm(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
+    elif args.model == 'Transformer':
+        model = Poetry_Model_Transformer(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
+        args.save_path='save_models/transformer.pth'
+    elif args.model == 'GPT-2':
+        model = Poetry_Model_lstm(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
+    else:
+        print("Please choose a model!\n")
+
+    model.load_state_dict(torch.load(args.save_path))
+    model = make_cuda(model)
+    poem = generate_poetry(model, string, w1, word_2_index, index_2_word,args)
+    print(poem)

src/apis/train.py

@@ -0,0 +1,68 @@
+import math
+import torch
+import numpy as np
+import torch.nn as nn
+import torch.optim as optim
+from src.utils.utils import make_cuda
+from torch.nn import functional as F
+from sklearn.metrics import mean_squared_error, mean_absolute_error
+
+
+def train(args, model, data_loader):
+    optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
+
+    model.train()
+    num_epochs = args.num_epochs
+
+    for epoch in range(num_epochs):
+        loss = 0
+        for step, (features, targets) in enumerate(data_loader):
+            features = make_cuda(features)
+            targets = make_cuda(targets)
+
+            optimizer.zero_grad()
+
+            pre, _ = model(features)
+            crs_loss = model.cross_entropy(pre, targets.reshape(-1))
+            loss += crs_loss.item()
+            crs_loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+            optimizer.step()
+
+            # print step info
+            if (step + 1) % args.log_step == 0:
+                print("Epoch [%.3d/%.3d] Step [%.3d/%.3d]: CROSS_loss=%.4f, RCROSS_loss=%.4f"
+                      % (epoch + 1,
+                         num_epochs,
+                         step + 1,
+                         len(data_loader),
+                         loss / args.log_step,
+                         math.sqrt(loss / args.log_step)))
+                loss = 0
+
+        # Loss = []
+        # for step, (features, targets) in enumerate(valid_data_loader):
+        #     features = make_cuda(features)
+        #     targets = make_cuda(targets)
+        #     model.eval()
+        #     preds = model(features)
+        #     valid_loss = CrossLoss(preds, targets)
+        #     Loss.append(valid_loss)
+        # print("Valid loss: %.3d\n" % (np.mean(Loss)))
+
+    return model
+
+
+def evaluate(args, model, data_loader):
+    model.eval()
+    loss = []
+    for step, (features, targets) in enumerate(data_loader):
+        features = make_cuda(features)
+        targets = make_cuda(targets)
+
+        pre, _ = model(features)
+        crs_loss = model.cross_entropy(pre, targets.reshape(-1))
+        loss.append(crs_loss.item())
+        torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+
+    print("loss=%.4f" % (np.mean(loss)))

src/datasets/dataloader.py

@@ -0,0 +1,115 @@
+import numpy as np
+import pickle
+import os
+import torch
+import torch.nn as nn
+from gensim.models.word2vec import Word2Vec
+from torch.utils.data import Dataset
+
+
+def padding(poetries, maxlen, pad):
+    batch_seq = [poetry + pad * (maxlen - len(poetry)) for poetry in poetries]
+    return batch_seq
+
+
+# 输入向后滑一字符为target，即预测下一个字
+def split_input_target(seq):
+    inputs = seq[:-1]
+    targets = seq[1:]
+    return inputs, targets
+
+
+# 创建词汇表
+def get_poetry(arg):
+    poetrys = []
+    if arg.Augmented_dataset:
+        path = arg.Augmented_data
+    else:
+        path = arg.data # 数据集路径,data/poetry.txt
+    with open(path, "r", encoding='UTF-8') as f:
+        for line in f:
+            try:
+                # line = line.decode('UTF-8')
+                line = line.strip(u'\n') # 去掉换行符
+                if arg.Augmented_dataset:
+                    content = line.strip(u' ')
+                else:
+                    title, content = line.strip(u' ').split(u':') # 标题和内容以冒号分隔
+                content = content.replace(u' ', u'') # 去掉空格
+                if u'_' in content or u'(' in content or u'（' in content or u'《' in content or u'[' in content: # 去掉特殊符号的古诗
+                    continue
+                if arg.strict_dataset: # 严格模式
+                    if len(content) < 12 or len(content) > 79:
+                        continue
+                else:
+                    if len(content) < 5 or len(content) > 79:
+                        continue
+                content = u'[' + content + u']' # 开头加上开始符，结尾加上结束符
+                poetrys.append(content) # 保存到poetrys列表中
+            except Exception as e:
+                pass
+
+            # 按诗的字数排序
+    poetrys = sorted(poetrys, key=lambda line: len(line))
+
+    with open("data/org_poetry.txt", "w", encoding="utf-8") as f:
+        for poetry in poetrys:
+            poetry = str(poetry).strip('[').strip(']').replace(',', '').replace('\'', '') + '\n'
+            f.write(poetry)
+
+    return poetrys
+
+
+# 切分文档
+def split_text(poetrys):
+    with open("data/split_poetry.txt", "w", encoding="utf-8") as f:
+        for poetry in poetrys:
+            poetry = str(poetry).strip('[').strip(']').replace(',', '').replace('\'', '') + '\n '
+            split_data = " ".join(poetry)
+            f.write(split_data)
+    return open("data/split_poetry.txt", "r", encoding='UTF-8').read()
+
+
+# 训练词向量
+def train_vec(split_file="data/split_poetry.txt", org_file="data/org_poetry.txt"):
+    param_file = "data/word_vec.pkl"
+    org_data = open(org_file, "r", encoding="utf-8").read().split("\n")
+    if os.path.exists(split_file):
+        all_data_split = open(split_file, "r", encoding="utf-8").read().split("\n")
+    else:
+        all_data_split = split_text().split("\n")
+
+    if os.path.exists(param_file):
+        return org_data, pickle.load(open(param_file, "rb"))
+
+    models = Word2Vec(all_data_split, vector_size=256, workers=7, min_count=1) # 训练词向量，输入参数分别是：分词后的文本，词向量维度，线程数，最小词频
+    pickle.dump([models.syn1neg, models.wv.key_to_index, models.wv.index_to_key], open(param_file, "wb")) # 保存词向量，key_to_index是词汇表，index_to_key是词向量，dump的作用是将数据序列化到文件中
+    return org_data, (models.syn1neg, models.wv.key_to_index, models.wv.index_to_key) # syn1neg是词向量，key_to_index是词汇表，index_to_key是词向量
+
+
+class Poetry_Dataset(Dataset):
+    def __init__(self, w1, word_2_index, all_data, Word2Vec):
+        self.Word2Vec = Word2Vec
+        self.w1 = w1
+        self.word_2_index = word_2_index
+        word_size, embedding_num = w1.shape
+        self.embedding = nn.Embedding(word_size, embedding_num) # 词嵌入层
+        # 最长句子长度
+        maxlen = max([len(seq) for seq in all_data])
+        pad = ' '
+        self.all_data = padding(all_data[:-1], maxlen, pad) 
+
+    def __getitem__(self, index):
+        a_poetry = self.all_data[index]
+
+        a_poetry_index = [self.word_2_index[i] for i in a_poetry]
+        xs, ys = split_input_target(a_poetry_index)
+        if self.Word2Vec:
+            xs_embedding = self.w1[xs]
+        else:
+            xs_embedding = np.array(xs)
+
+        return xs_embedding, np.array(ys).astype(np.int64)
+
+    def __len__(self):
+        return len(self.all_data)

src/models/LSTM/model.py

@@ -0,0 +1,37 @@
+import torch
+import numpy as np
+import torch.nn as nn
+
+
+class Poetry_Model_lstm(nn.Module):
+    def __init__(self, hidden_num, word_size, embedding_num, Word2Vec):
+        super().__init__()
+
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.hidden_num = hidden_num
+        self.Word2Vec = Word2Vec
+
+        self.embedding = nn.Embedding(word_size, embedding_num)
+        self.lstm = nn.LSTM(input_size=embedding_num, hidden_size=hidden_num, batch_first=True, num_layers=2,
+                            bidirectional=False)
+        self.dropout = nn.Dropout(0.3)
+        self.flatten = nn.Flatten(0, 1)
+        self.linear = nn.Linear(hidden_num, word_size)
+        self.cross_entropy = nn.CrossEntropyLoss()
+
+    def forward(self, xs_embedding, h_0=None, c_0=None):
+        # xs_embedding: [batch_size, max_seq_len, n_feature] n_feature=256
+        if h_0 == None or c_0 == None:
+            h_0 = torch.tensor(np.zeros((2, xs_embedding.shape[0], self.hidden_num), dtype=np.float32))
+            c_0 = torch.tensor(np.zeros((2, xs_embedding.shape[0], self.hidden_num), dtype=np.float32))
+        h_0 = h_0.to(self.device)
+        c_0 = c_0.to(self.device)
+        xs_embedding = xs_embedding.to(self.device)
+        if not self.Word2Vec:
+            xs_embedding = self.embedding(xs_embedding)
+        hidden, (h_0, c_0) = self.lstm(xs_embedding, (h_0, c_0))
+        hidden_drop = self.dropout(hidden)
+        hidden_flatten = self.flatten(hidden_drop)
+        pre = self.linear(hidden_flatten)
+        # pre：[batch_size*max_seq_len, vocab_size]
+        return pre, (h_0, c_0)

src/models/Transformer/model.py

@@ -0,0 +1,70 @@
+import torch
+import numpy as np
+import torch.nn as nn
+import math
+
+class Poetry_Model_Transformer(nn.Module):
+    def __init__(self, hidden_num, word_size, embedding_num, Word2Vec):
+        super().__init__()
+
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.hidden_num = hidden_num
+        self.Word2Vec = Word2Vec
+
+        # 位置编码
+        self.pos_encoder= PositionalEncoding(d_model=embedding_num,dropout=0.5)
+        self.embedding = nn.Embedding(word_size, embedding_num) 
+        self.transformer = nn.Transformer(d_model=embedding_num, nhead=8, num_encoder_layers=6, num_decoder_layers=6,
+                                            dim_feedforward=2048, dropout=0.5, activation='relu') # 输入的参数分别是：词嵌入的维度，多头注意力的头数，编码器层数，解码器层数，前馈网络的隐藏层维度，dropout概率，激活函数
+        
+        # 编码器
+        self.encoder_layer=nn.TransformerEncoderLayer(d_model=embedding_num,nhead=8,dim_feedforward=2048,dropout=0.5)
+        self.encoder_norm=nn.LayerNorm(normalized_shape=embedding_num)
+        self.encoder=nn.TransformerEncoder(self.encoder_layer,num_layers=6,norm=self.encoder_norm)
+
+        # 解码器
+        # self.decoder_layer=nn.TransformerDecoderLayer(d_model=embedding_num,nhead=8,dim_feedforward=2048,dropout=0.5)
+        # self.decoder_norm=nn.LayerNorm(normalized_shape=embedding_num)
+        # self.decoder=nn.TransformerDecoder(self.decoder_layer,num_layers=6,norm=self.decoder_norm)
+        self.flatten = nn.Flatten(0, 1)
+        self.linear1 = nn.Linear(embedding_num, hidden_num)
+        self.linear2 = nn.Linear(hidden_num, word_size)
+        self.cross_entropy = nn.CrossEntropyLoss()
+
+    def forward(self, xs_embedding, h_0=None, c_0=None):
+        if h_0 == None or c_0 == None:
+            h_0 = torch.tensor(np.zeros((2, xs_embedding.shape[0], self.hidden_num), dtype=np.float32))
+            c_0 = torch.tensor(np.zeros((2, xs_embedding.shape[0], self.hidden_num), dtype=np.float32))
+        h_0 = h_0.to(self.device)
+        c_0 = c_0.to(self.device)
+        xs_embedding = xs_embedding.to(self.device)
+
+        if not self.Word2Vec:
+            xs_embedding = self.embedding(xs_embedding)
+
+        encoder_input = self.pos_encoder(xs_embedding)
+        pre_encoded=self.encoder(encoder_input)
+        pre=self.linear2(self.linear1(self.flatten(pre_encoded)))
+        
+        # pre：[batch_size*max_seq_len, vocab_size]
+        return pre, (h_0, c_0)
+    
+class PositionalEncoding(nn.Module):
+
+    def __init__(self, d_model, dropout = 0.1, max_len = 5000):
+        super(PositionalEncoding, self).__init__()
+
+        self.dropout = nn.Dropout(p = dropout)
+
+        pe = torch.zeros(max_len, d_model)
+
+        position = torch.arange(0, max_len, dtype = torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        x = x + self.pe[:x.size(0), :]
+        return self.dropout(x)

src/utils/utils.py

@@ -0,0 +1,15 @@
+import torch
+
+
+def make_cuda(tensor):
+    """Use CUDA if it's available."""
+    if torch.cuda.is_available():
+        tensor = tensor.cuda()
+    return tensor
+
+
+def is_minimum(value, indiv_to_rmse):
+    if len(indiv_to_rmse) == 0:
+        return True
+    temp = list(indiv_to_rmse.values())
+    return True if value < min(temp) else False

train.py

@@ -0,0 +1,70 @@
+from src.utils.utils import make_cuda
+from src.apis.train import train, evaluate
+from src.models.LSTM.model import Poetry_Model_lstm
+import argparse
+import torch
+import os
+from src.datasets.dataloader import Poetry_Dataset, train_vec, get_poetry, split_text
+from torch.utils.data import DataLoader
+from src.models.Transformer.model import Poetry_Model_Transformer
+
+
+def parse_arguments():
+    # argument parsing
+    parser = argparse.ArgumentParser(description="Specify Params for Experimental Setting")
+
+    parser.add_argument('--batch_size', type=int, default=64,
+                        help="Specify batch size")
+
+    parser.add_argument('--initial_epochs', type=int, default=25,
+                        help="Specify the number of epochs for initial training")
+
+    parser.add_argument('--num_epochs', type=int, default=50,
+                        help="Specify the number of epochs for competitive search")
+    parser.add_argument('--log_step', type=int, default=100,
+                        help="Specify log step size for training")
+    parser.add_argument('--learning_rate', type=float, default=1e-3,
+                        help="Learning rate")
+    parser.add_argument('--data', type=str, default='data/poetry.txt',
+                        help="Path to the dataset")
+    parser.add_argument('--n_hidden', type=int, default=128)
+    parser.add_argument('--max_grad_norm', type=float, default=1.0)
+
+    parser.add_argument('--save_path', type=str, default='save_models/transformer.pth')
+    parser.add_argument('--strict_dataset', default=False, help="strict dataset")
+    parser.add_argument('--Word2Vec',type=bool, default=True)
+    parser.add_argument("--Augmented_dataset", type=bool, default=False)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_arguments()
+    if os.path.exists("data/split_poetry.txt") and os.path.exists("data/org_poetry.txt"):
+        print("pre_file exit!")
+    else:
+        split_text(get_poetry(args)) # split poetry
+    all_data, (w1, word_2_index, index_2_word) = train_vec()
+    args.word_size, args.embedding_num = w1.shape # 词向量的维度
+
+    dataset = Poetry_Dataset(w1, word_2_index, all_data, Word2Vec=args.Word2Vec)
+    train_size = int(len(dataset) * 0.7)
+    test_size = len(dataset) - train_size
+    train_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size])
+
+    train_data_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
+    valid_data_loader = DataLoader(test_dataset, batch_size=int(args.batch_size/4), shuffle=True)
+    
+    # best_model = Poetry_Model_lstm(args.n_hidden, args.word_size, args.embedding_num,args.Word2Vec)
+    best_model = Poetry_Model_Transformer(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
+    best_model = make_cuda(best_model) # use gpu
+    print("Initial training before competitive random search")
+    best_model = train(args, best_model, train_data_loader)
+
+    torch.save(best_model.state_dict(), args.save_path)
+
+    print('test evaluation:')
+    evaluate(args, best_model, valid_data_loader)
+
+
+if __name__ == '__main__':
+    main()