change train py some

config.py

@@ -24,7 +24,7 @@ class Config(object):
         self.num_epochs = 3
         self.log_batch = 100
         self.batch_size = 128
-        self.max_seq_len = 32
+        self.max_seq_len = 128
         self.require_improvement = 1000
 
         self.warmup_steps = 0

main.py

@@ -12,6 +12,9 @@ from config import Config
 from preprocess import DataProcessor, get_time_dif
 from transformers import BertConfig, BertTokenizer, BertForSequenceClassification
 
+from ebart import PegasusSummarizer
+
+
 parser = argparse.ArgumentParser(description="Bert Chinese Text Classification")
 parser.add_argument("--mode", type=str, required=True, help="train/demo/predict")
 parser.add_argument("--data_dir", type=str, default="./data", help="training data and saved model path")
@@ -40,11 +43,14 @@ def main():
     )
     model.to(config.device)
 
+    #
+    summarizerModel = PegasusSummarizer()
+
     if args.mode == "train":
         print("loading data...")
         start_time = time.time()
-        train_iterator = DataProcessor(config.train_file, config.device, tokenizer, config.batch_size, config.max_seq_len, args.seed)
-        dev_iterator = DataProcessor(config.dev_file, config.device, tokenizer, config.batch_size, config.max_seq_len, args.seed)
+        train_iterator = DataProcessor(config.train_file, config.device, summarizerModel,tokenizer, config.batch_size, config.max_seq_len, args.seed)
+        dev_iterator = DataProcessor(config.dev_file, config.device, summarizerModel,tokenizer, config.batch_size, config.max_seq_len, args.seed)
         time_dif = get_time_dif(start_time)
         print("time usage:", time_dif)
 

preprocess.py

@@ -6,6 +6,7 @@ import torch
 import random
 from tqdm import tqdm
 from datetime import timedelta
+from ebart import PegasusSummarizer
 
 def get_time_dif(start_time):
     end_time = time.time()
@@ -14,10 +15,11 @@ def get_time_dif(start_time):
 
 ## 获取张量数据
 class DataProcessor(object):
-    def __init__(self, dataPath, device, tokenizer, batch_size, max_seq_len, seed):
+    def __init__(self, dataPath, device,summarizerModel, tokenizer, batch_size, max_seq_len, seed):
         self.seed = seed
         self.device = device
         self.tokenizer = tokenizer
+        self.summarizerModel=summarizerModel
         self.batch_size = batch_size
         self.max_seq_len = max_seq_len
 
@@ -51,6 +53,7 @@ class DataProcessor(object):
                 content, label = line.rsplit('\t', 1)
                 mapped_label = labels_map.get(label)
                 if mapped_label is not None and isinstance(mapped_label, int) and mapped_label >= 0:
+                    content= self.summarizerModel.generate_summary(content, 128, 64)
                     contents.append(content)
                     labels.append(mapped_label)
                 else: