add some format

Author: Lan-ping

.gitignore

@@ -129,3 +129,5 @@ dmypy.json
 .pyre/
 
 chinese-roberta-wwm-ext/
+
+output/*.pth

data.py

@@ -17,8 +17,12 @@ def __getitem__(self, item):
 
         encoding = self.tokenizer.encode_plus(
             review,
+            add_special_tokens=True,
             max_length=self.max_len,
+            return_token_type_ids=True,
             padding='max_length',
+            truncation=True,
+            return_attention_mask=True,
             return_tensors='pt')
 
         return {

main.py

@@ -7,6 +7,8 @@
 from model import TextBackbone
 from data import ReviewsDataset, load_data
 
+from transformers import AdamW
+
 def parse_arguments():
     parser = argparse.ArgumentParser(
         description='Simple Sentiment Analysis with PyTorch and Transformers'
@@ -15,7 +17,7 @@ def parse_arguments():
     parser.add_argument('--n_classes', default=2, type=int, help='number of classes')
 
     parser.add_argument('--data_path', type=str, default='data/data.txt', help='the path of dataset')
-    parser.add_argument('--batch_size', default=20, type=int, help='batch size')
+    parser.add_argument('--batch_size', default=8, type=int, help='batch size')
 
     parser.add_argument('--epochs', default=50, type=int, help='number of epochs tp train for')
     parser.add_argument('--lr', default=1e-3, type=float, help='learning rate')
@@ -28,9 +30,12 @@ def parse_arguments():
 def train(model, dataset, optimizer, device, batch_size, epochs):
     model.train()
     train_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
-    pbar = tqdm(train_loader)
+    
+    min_loss = float('inf')
     for epoch in range(epochs):
+        pbar = tqdm(train_loader)
         pbar.set_description("Epoch {}:".format(epoch))
+        total_loss = 0
         for batch in pbar:
             batch = {key: value.to(device) for key, value in batch.items()}
             optimizer.zero_grad()
@@ -39,9 +44,13 @@ def train(model, dataset, optimizer, device, batch_size, epochs):
             loss.backward()
             optimizer.step()
             pbar.set_postfix(loss=loss.item())
+            total_loss += loss.item()
 
-        if epoch % 10 == 0:
-            model.save('model_{}.pth'.format(epoch))
+        if total_loss < min_loss:
+            min_loss = total_loss
+            torch.save(model.state_dict(), 'output/model_best.pth')
+
+        print("Epoch {}: Average loss: {}".format(epoch, total_loss / len(train_loader)))
 
     return model
 
@@ -51,7 +60,8 @@ def main():
     reviews, targets = load_data(args.data_path)
     dataset = ReviewsDataset(reviews, targets)
     model = TextBackbone(num_classes=args.n_classes).to(args.device)
-    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+    optimizer = AdamW(model.parameters(),lr=2e-5, eps=1e-8)
+
     model = train(model, dataset, optimizer, args.device, args.batch_size, args.epochs)
     torch.save(model.state_dict(), 'model.pth')
 

output/README.md

@@ -0,0 +1 @@
+# Output