Upload 53 files

.idea/.gitignore

@@ -0,0 +1,3 @@
+# Default ignored files
+/shelf/
+/workspace.xml

.idea/.name

@@ -0,0 +1 @@
+inference.py

.idea/ea_lstm.iml

@@ -0,0 +1,8 @@
+<?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>
+</module>

.idea/inspectionProfiles/Project_Default.xml

@@ -0,0 +1,20 @@
+<component name="InspectionProjectProfileManager">
+  <profile version="1.0">
+    <option name="myName" value="Project Default" />
+    <inspection_tool class="PyPackageRequirementsInspection" enabled="true" level="WARNING" enabled_by_default="true">
+      <option name="ignoredPackages">
+        <value>
+          <list size="7">
+            <item index="0" class="java.lang.String" itemvalue="easydict" />
+            <item index="1" class="java.lang.String" itemvalue="pandas" />
+            <item index="2" class="java.lang.String" itemvalue="matplotlib" />
+            <item index="3" class="java.lang.String" itemvalue="pillow" />
+            <item index="4" class="java.lang.String" itemvalue="mindspore" />
+            <item index="5" class="java.lang.String" itemvalue="setuptools" />
+            <item index="6" class="java.lang.String" itemvalue="numpy" />
+          </list>
+        </value>
+      </option>
+    </inspection_tool>
+  </profile>
+</component>

.idea/misc.xml

@@ -1,4 +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" />
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.9 (pytorch)" project-jdk-type="Python SDK" />
 </project>

.idea/modules.xml

@@ -2,7 +2,7 @@
 <project version="4">
   <component name="ProjectModuleManager">
     <modules>
-      <module fileurl="file://$PROJECT_DIR$/.idea/nlp.iml" filepath="$PROJECT_DIR$/.idea/nlp.iml" />
+      <module fileurl="file://$PROJECT_DIR$/.idea/ea_lstm.iml" filepath="$PROJECT_DIR$/.idea/ea_lstm.iml" />
     </modules>
   </component>
 </project>

.idea/workspace.xml

@@ -1,111 +1,44 @@
 <?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>
+    <list default="true" id="276a53df-3cdd-4e96-95d3-c1e69d4e9b9f" name="Changes" comment="" />
     <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="ProjectId" id="2OyFWrJQpFYHFKgf87OgmRH5Jtu" />
   <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 name="PropertiesComponent"><![CDATA[{
+  "keyToString": {
+    "RunOnceActivity.OpenProjectViewOnStart": "true",
+    "RunOnceActivity.ShowReadmeOnStart": "true",
+    "last_opened_file_path": "C:/Users/LENOVO/PycharmProjects/lstm"
   }
-}</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>
   <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>
+      <changelist id="276a53df-3cdd-4e96-95d3-c1e69d4e9b9f" name="Changes" comment="" />
+      <created>1682524950142</created>
       <option name="number" value="Default" />
       <option name="presentableId" value="Default" />
-      <updated>1684726163448</updated>
+      <updated>1682524950142</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 name="XDebuggerManager">
+    <watches-manager>
+      <configuration name="PythonConfigurationType">
+        <watch expression="input_eval" />
+        <watch expression="word_2_index" />
+      </configuration>
+    </watches-manager>
   </component>
 </project>

README.md

@@ -1,12 +1,23 @@
----
-title: Poem Generation
-emoji: 👁
-colorFrom: gray
-colorTo: red
-sdk: gradio
-sdk_version: 3.32.0
-app_file: app.py
-pinned: false
----
+# NLP Final Project
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+```shell
+├── configs
+├── data
+│   └── poetry.txt
+├── inference.py
+├── src
+│   ├── apis
+│   │   ├── evaluate.py
+│   │   ├── inference.py
+│   │   └── train.py
+│   ├── datasets
+│   │   └── dataloader.py
+│   ├── models
+│   │   └── EA-LSTM
+│   │       ├── algorithm.py
+│   │       └── model.py
+│   └── utils
+│       └── utils.py
+├── test.py
+└── train.py
+```

app.py

@@ -7,7 +7,7 @@ from inference import infer
 
 
 
-INTERFACE = gradio.Interface(fn=infer, inputs=[gradio.Radio(["lstm","GRU","Seq2Seq","Transformer","GPT-2"]),"text"], outputs=["text"], title="Poetry Generation",
+INTERFACE = gradio.Interface(fn=infer, inputs=[gradio.Radio(["lstm","GRU"]),"text"], outputs=["text"], title="Poetry Generation",
                              description="Choose a model and input the poetic head to generate a acrostic",
                              thumbnail="https://github.com/gradio-app/gpt-2/raw/master/screenshots/interface.png?raw=true")
 

data/word_vec.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1164cfc2e28ef6ecbb1a04734e7268238b4841667f13d6cb4c42e27717dd4575
+size 6339344

inference.py

@@ -1,10 +1,10 @@
 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
-from src.models.LSTM.model import Poetry_Model_lstm
+
 
 def parse_arguments():
     # argument parsing
@@ -15,12 +15,12 @@ def parse_arguments():
     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')
+    parser.add_argument('--save_path', type=str, default='save_models/lstm_50.pth')
 
     return parser.parse_args()
 
 
-def generate_poetry(model, head_string, w1, word_2_index, index_2_word,args):
+def generate_poetry(model, head_string, w1, word_2_index, index_2_word):
     print("藏头诗生成中...., {}".format(head_string))
     poem = ""
     # 以句子的每一个字为开头生成诗句
@@ -54,33 +54,31 @@ def generate_poetry(model, head_string, w1, word_2_index, index_2_word,args):
 
     return poem
 
-def infer(model,poem_head):
+def infer(model,string):
     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 = input("诗头:")
     # string = '自然语言'
-
+    args.model=model
     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'
+        args.save_path = 'save_models/lstm_50.pth'
     elif args.model == 'GRU':
         model = Poetry_Model_lstm(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
+        args.save_path = 'save_models/GRU_50.pth'
     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_100.pth'
+        model = Poetry_Model_lstm(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
     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,map_location=torch.device('cpu')))
+    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)
+    poem = generate_poetry(model, string, w1, word_2_index, index_2_word)
     return poem
 
 
@@ -88,25 +86,23 @@ 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 = '自然语言'
-    
+    # 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_100.pth'
+        model = Poetry_Model_lstm(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
     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,map_location=torch.device('cpu')))
+    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)
+    poem = generate_poetry(model, string, w1, word_2_index, index_2_word)
     print(poem)

save_models/GRU_25.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bacf9a7ec329c6185098c1309ab28239b4c087b53832b3d18e5323831bfead23
+size 10727391

save_models/GRU_50.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9a8e83a733c023b35c44020e014bb72e2c1d05698eb782669c0e4d5a76d4590d
+size 10727391

save_models/lstm_25.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1501b5e3e6d9aa864857c8036f27d09c2489da832e616916b9633092b0ed3df5
+oid sha256:2b064666ce02c63541dee4b6146d31ee8f7e784ee9c2811c9b9266aba6cc4193
 size 10727391

save_models/lstm_50.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa157d970149c32b53b024a23ef8428e7b7e1702ed72d44152b568b085b1bfaa
+size 10727391

src/apis/evaluate.py

@@ -0,0 +1,23 @@
+import torch
+import numpy as np
+from src.models.EA_LSTM.model import weightedLSTM
+from src.datasets.dataloader import MyDataset, create_vocab
+
+
+def test(args):
+    vocab, poetrys = create_vocab(args.data)
+    # 词汇表长度
+    args.vocab_size = len(vocab)
+    int2char = np.array(vocab)
+    valid_dataset = MyDataset(vocab, poetrys, args, train=False)
+
+    model = weightedLSTM(6110, 256, 128, 2, [1.0] * 80, False)
+    model.load_state_dict(torch.load(args.save_path))
+
+    input_example_batch, target_example_batch = valid_dataset[0]
+    example_batch_predictions = model(input_example_batch)
+    predicted_id = torch.distributions.Categorical(example_batch_predictions).sample()
+    predicted_id = torch.squeeze(predicted_id, -1).numpy()
+    print("Input: \n", repr("".join(int2char[input_example_batch])))
+    print()
+    print("Predictions: \n", repr("".join(int2char[predicted_id])))

src/apis/train.py

@@ -8,11 +8,11 @@ from torch.nn import functional as F
 from sklearn.metrics import mean_squared_error, mean_absolute_error
 
 
-def train(args, model, data_loader):
+def train(args, model, data_loader, initial=False):
     optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
 
     model.train()
-    num_epochs = args.num_epochs
+    num_epochs = args.initial_epochs if initial else args.num_epochs
 
     for epoch in range(num_epochs):
         loss = 0

src/datasets/dataloader.py

@@ -25,27 +25,27 @@ def get_poetry(arg):
     if arg.Augmented_dataset:
         path = arg.Augmented_data
     else:
-        path = arg.data # 数据集路径,data/poetry.txt
+        path = arg.data
     with open(path, "r", encoding='UTF-8') as f:
         for line in f:
             try:
                 # line = line.decode('UTF-8')
-                line = line.strip(u'\n') # 去掉换行符
+                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: # 去掉特殊符号的古诗
+                    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 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列表中
+                content = u'[' + content + u']'
+                poetrys.append(content)
             except Exception as e:
                 pass
 
@@ -82,9 +82,9 @@ def train_vec(split_file="data/split_poetry.txt", org_file="data/org_poetry.txt"
     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是词向量
+    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"))
+    return org_data, (models.syn1neg, models.wv.key_to_index, models.wv.index_to_key)
 
 
 class Poetry_Dataset(Dataset):
@@ -93,11 +93,11 @@ class Poetry_Dataset(Dataset):
         self.w1 = w1
         self.word_2_index = word_2_index
         word_size, embedding_num = w1.shape
-        self.embedding = nn.Embedding(word_size, embedding_num) # 词嵌入层
+        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) 
+        self.all_data = padding(all_data[:-1], maxlen, pad)
 
     def __getitem__(self, index):
         a_poetry = self.all_data[index]

src/models/LSTM/model.py

@@ -20,7 +20,7 @@ class Poetry_Model_lstm(nn.Module):
         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
+        # xs_embedding: [batch_size, max_seq_len, n_feature] n_feature=128
         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))

train.py

@@ -6,19 +6,20 @@ 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('--model', type=str, default='lstm',
+                        help="lstm/GRU/Seq2Seq/Transformer/GPT-2")
+    parser.add_argument('--Word2Vec', default=True)
+    parser.add_argument('--Augmented_dataset', default=False, help="augmented dataset")
+    parser.add_argument('--strict_dataset', default=False, help="strict dataset")
+
     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,
@@ -27,40 +28,51 @@ def parse_arguments():
                         help="Learning rate")
     parser.add_argument('--data', type=str, default='data/poetry.txt',
                         help="Path to the dataset")
+    parser.add_argument('--Augmented_data', type=str, default='data/poetry_7.txt',
+                        help="Path to the Augmented_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)
+    parser.add_argument('--save_path', type=str, default='save_models/')
+
     return parser.parse_args()
 
 
 def main():
     args = parse_arguments()
+    # if you want to change the data(org data or argument data), please delete file: 'split_poetry.txt' and 'org_poetry.txt'
     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
+        split_text(get_poetry(args))
+
     all_data, (w1, word_2_index, index_2_word) = train_vec()
-    args.word_size, args.embedding_num = w1.shape # 词向量的维度
+    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)
+    dataset = Poetry_Dataset(w1, word_2_index, all_data, args.Word2Vec)
+    train_size = int(len(dataset) * 0.8)
     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")
+
+    if args.model == 'lstm':
+        best_model = Poetry_Model_lstm(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
+    elif args.model == 'GRU':
+        best_model = Poetry_Model_lstm(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
+    elif args.model == 'Seq2Seq':
+        best_model = Poetry_Model_lstm(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
+    elif args.model == 'Transformer':
+        best_model = Poetry_Model_lstm(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
+    elif args.model == 'GPT-2':
+        best_model = Poetry_Model_lstm(args.n_hidden, args.word_size, args.embedding_num, args.Word2Vec)
+    else:
+        print("Please choose a model!\n")
+
+    best_model = make_cuda(best_model)
     best_model = train(args, best_model, train_data_loader)
 
-    torch.save(best_model.state_dict(), args.save_path)
+    torch.save(best_model.state_dict(), args.save_path + args.model + '_' + str(args.num_epochs)+'.pth')
 
     print('test evaluation:')
     evaluate(args, best_model, valid_data_loader)