contrastive commit 3

data/{twitter-unsup.csv → amazon-polarity.parquet}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5a7af1ec5fc749ec8e5ea13c574aeb5c06254aa1c081e3421868079d5356b3f4
-size 20895533
+oid sha256:dbe4770cfa6be45add6c9a322044bd4c1901520dde5a2707eca402a74fbe854e
+size 870289

unsup_simcse.py

@@ -3,6 +3,7 @@ import torch
 import random
 import argparse
 import numpy as np
+import pandas as pd
 import torch.nn.functional as F
 
 from tqdm import tqdm
@@ -20,7 +21,7 @@ from classifier import SentimentDataset, BertSentimentClassifier
 TQDM_DISABLE = False
 
 
-class TwitterDataset(Dataset):
+class AmazonDataset(Dataset):
     def __init__(self, dataset, args):
         self.dataset = dataset
         self.p = args
@@ -31,19 +32,22 @@ class TwitterDataset(Dataset):
     def __getitem__(self, idx):
         return self.dataset[idx]
     
-    def pad_data(self, sents):
+    def pad_data(self, data):
+        sents = [x[0] for x in data]
+        sent_ids = [x[1] for x in data]
         encoding = tokenizer(sents, return_tensors='pt', padding=True, truncation=True)
         token_ids = torch.LongTensor(encoding['input_ids'])
         attension_mask = torch.LongTensor(encoding['attention_mask'])
 
-        return token_ids, attension_mask
+        return token_ids, attension_mask, sent_ids
     
-    def collate_fn(self, sents):
-        token_ids, attention_mask = self.pad_data(sents)
+    def collate_fn(self, data):
+        token_ids, attention_mask, sent_ids = self.pad_data(data)
 
         batched_data = {
             'token_ids': token_ids,
             'attention_mask': attention_mask,
+            'sent_ids': sent_ids
         }
 
         return batched_data
@@ -51,36 +55,36 @@ class TwitterDataset(Dataset):
 
 def load_data(filename, flag='train'):
     '''
-    - for Twitter dataset: list of sentences
-    - for SST/CFIMDB dataset: list of (sent, [label], sent_id)
+    - for amazon dataset: list of (sent, sent_id)
+    - for test dataset: list of (sent, sent_id)
+    - for train dataset: list of (sent, label, sent_id)
     '''
 
-    num_labels = set()
-    data = []
-    with open(filename, 'r') as fp:
-        if flag == 'twitter':
-            for cnt, record in enumerate(csv.DictReader(fp, delimiter = ',')):
-                sent = record['clean_text'].lower().strip()
-                data.append(sent)
-                if cnt == 10000: break
-        elif flag == 'test':
-            for record in csv.DictReader(fp, delimiter = '\t'):
-                sent = record['sentence'].lower().strip()
-                sent_id = record['id'].lower().strip()
-                data.append((sent,sent_id))
-        else:
-            for record in csv.DictReader(fp, delimiter = '\t'):
-                sent = record['sentence'].lower().strip()
-                sent_id = record['id'].lower().strip()
-                label = int(record['sentiment'].strip())
-                num_labels.add(label)
-                data.append((sent, label, sent_id))
-        print(f"load {len(data)} data from {filename}")
-
-    if flag == 'train':
-        return data, len(num_labels)
+    if flag == 'amazon':
+        df = pd.read_parquet(filename)
+        data = list(zip(df['content'], df.index))
     else:
+        data, num_labels = [], set()
+
+        with open(filename, 'r') as fp:
+            if flag == 'test':
+                for record in csv.DictReader(fp, delimiter = '\t'):
+                    sent = record['sentence'].lower().strip()
+                    sent_id = record['id'].lower().strip()
+                    data.append((sent,sent_id))
+            else:
+                for record in csv.DictReader(fp, delimiter = '\t'):
+                    sent = record['sentence'].lower().strip()
+                    sent_id = record['id'].lower().strip()
+                    label = int(record['sentiment'].strip())
+                    num_labels.add(label)
+                    data.append((sent, label, sent_id))
+
+    print(f"load {len(data)} data from {filename}")
+    if flag in ['test', 'amazon']:
         return data
+    else:
+        return data, len(num_labels)
 
 
 def save_model(model, optimizer, args, config, filepath):
@@ -98,11 +102,6 @@ def save_model(model, optimizer, args, config, filepath):
     print(f"save the model to {filepath}")
 
 
-# def model_eval(dataloader, model, device):
-#     model.eval()
-
-
-
 def contrastive_loss(embeds_1: Tensor, embeds_2: Tensor, temp=0.05):
     '''
     embeds_1: [batch_size, hidden_size]
@@ -131,7 +130,7 @@ def train(args):
     '''
     Training Pipeline
     -----------------
-    1. Load the Twitter Sentiment and SST Dataset.
+    1. Load the Amazon Polarity and SST Dataset.
     2. Determine batch_size (64) and number of batches (?).
     3. Initialize SentimentClassifier (including bert).
     4. Looping through 10 epoches.
@@ -142,16 +141,16 @@ def train(args):
     9. If dev_acc > best_dev_acc: save_model(...)
     '''
 
-    twitter_data = load_data(args.train_bert, 'twitter')
+    amazon_data = load_data(args.train_bert, 'amazon')
     train_data, num_labels = load_data(args.train, 'train')
     dev_data = load_data(args.dev, 'valid')
 
-    twitter_dataset = TwitterDataset(twitter_data, args)
+    amazon_dataset = AmazonDataset(amazon_data, args)
     train_dataset = SentimentDataset(train_data, args)
     dev_dataset = SentimentDataset(dev_data, args)
 
-    twitter_dataloader = DataLoader(twitter_dataset, shuffle=True, batch_size=args.batch_size_cse,
-                                    num_workers=args.num_cpu_cores, collate_fn=twitter_dataset.collate_fn)
+    amazon_dataloader = DataLoader(amazon_dataset, shuffle=True, batch_size=args.batch_size_cse,
+                                    num_workers=args.num_cpu_cores, collate_fn=amazon_dataset.collate_fn)
     train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=args.batch_size_classifier,
                                   num_workers=args.num_cpu_cores, collate_fn=train_dataset.collate_fn)
     dev_dataloader = DataLoader(dev_dataset, shuffle=False, batch_size=args.batch_size_classifier,
@@ -177,7 +176,7 @@ def train(args):
     for epoch in range(args.epochs):
         model.bert.train()
         train_loss = num_batches = 0
-        for batch in tqdm(twitter_dataloader, f'train-twitter-{epoch}', leave=False, disable=TQDM_DISABLE):
+        for batch in tqdm(amazon_dataloader, f'train-amazon-{epoch}', leave=False, disable=TQDM_DISABLE):
             b_ids, b_mask = batch['token_ids'], batch['attention_mask']
             b_ids = b_ids.to(device)
             b_mask = b_mask.to(device)
@@ -189,11 +188,13 @@ def train(args):
             # Calculate mean SimCSE loss function
             loss = contrastive_loss(logits_1, logits_2)
 
+            # Back propagation
+            optimizer_cse.zero_grad()
             loss.backward()
             optimizer_cse.step()
 
             train_loss += loss.item()
-            num_batches += 0
+            num_batches += 1
 
         train_loss = train_loss / num_batches
         print(f"Epoch {epoch}: train loss :: {train_loss :.3f}")
@@ -205,11 +206,12 @@ def get_args():
     parser.add_argument("--num-cpu-cores", type=int, default=4)
     parser.add_argument("--epochs", type=int, default=10)
     parser.add_argument("--use_gpu", action='store_true')
-    parser.add_argument("--batch_size_cse", help="'unsup': 64, 'sup': 512", type=int)
-    parser.add_argument("--batch_size_classifier", help="'sst': 64, 'cfimdb': 8", type=int)
+    parser.add_argument("--batch_size_cse", type=int, default=8)
+    parser.add_argument("--batch_size_sst", type=int, default=64)
+    parser.add_argument("--batch_size_cfimdb", type=int, default=8)
     parser.add_argument("--hidden_dropout_prob", type=float, default=0.3)
-    parser.add_argument("--lr_cse", default=2e-5)
-    parser.add_argument("--lr_classifier", default=1e-5)
+    parser.add_argument("--lr_cse", type=float, default=2e-5)
+    parser.add_argument("--lr_classifier", type=float, default=1e-5)
 
     args = parser.parse_args()
     return args
@@ -229,9 +231,9 @@ if __name__ == "__main__":
         use_gpu=args.use_gpu,
         epochs=args.epochs,
         batch_size_cse=args.batch_size_cse,
-        batch_size_classifier=args.batch_size_classifier,
+        batch_size_classifier=args.batch_size_sst,
         hidden_dropout_prob=args.hidden_dropout_prob,
-        train_bert='data/twitter-unsup.csv',
+        train_bert='data/amazon-polarity.parquet',
         train='data/ids-sst-train.csv',
         dev='data/ids-sst-dev.csv',
         test='data/ids-sst-test-student.csv'