contrastive commit 2

.gitignore

@@ -163,4 +163,5 @@ cython_debug/
 
 nohup.out
 *.pt
-predictions/
+predictions/
+zemo*.py

classifier.py

@@ -362,8 +362,6 @@ def main():
     seed_everything(args.seed)
     torch.set_num_threads(args.num_cpu_cores)
 
-    print(torch.get_num_threads())
-
     print('Training Sentiment Classifier on SST...')
     config = SimpleNamespace(
         filepath='sst-classifier.pt',

unsup_simcse.py

@@ -3,8 +3,10 @@ import torch
 import random
 import argparse
 import numpy as np
+import torch.nn.functional as F
 
 from tqdm import tqdm
+from torch import Tensor
 from types import SimpleNamespace
 from torch.utils.data import Dataset, DataLoader
 from sklearn.metrics import f1_score, accuracy_score
@@ -52,18 +54,22 @@ def load_data(filename, flag='train'):
     - for Twitter dataset: list of sentences
     - for SST/CFIMDB dataset: list of (sent, [label], sent_id)
     '''
+
     num_labels = set()
     data = []
     with open(filename, 'r') as fp:
-        for record in csv.DictReader(fp, delimiter = ',', ):
-            if flag == 'twitter':
+        if flag == 'twitter':
+            for cnt, record in enumerate(csv.DictReader(fp, delimiter = ',')):
                 sent = record['clean_text'].lower().strip()
                 data.append(sent)
-            elif flag == 'test':
+                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:
+        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())
@@ -92,6 +98,35 @@ 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]
+    embeds_2: [batch_size, hidden_size]
+    '''
+
+    # [batch_size, batch_size]
+    sim_matrix = F.cosine_similarity(embeds_1.unsqueeze(1), embeds_2.unsqueeze(0), dim=-1) / temp
+
+    # [batch_size]
+    positive_sim = torch.diagonal(sim_matrix)
+
+    # [batch_size]
+    nume = torch.exp(positive_sim)
+
+    # [batch_size]
+    deno = torch.exp(sim_matrix).sum(1)
+
+    # [batch_size]
+    loss_per_batch = -torch.log(nume / deno)
+
+    return loss_per_batch.mean()
+
+
 def train(args):
     '''
     Training Pipeline
@@ -138,6 +173,7 @@ def train(args):
     optimizer_classifier = AdamW(model.parameters(), lr=args.lr_classifier)
     best_dev_acc = 0
 
+    # ---- Training minBERT using SimCSE ---- #
     for epoch in range(args.epochs):
         model.bert.train()
         train_loss = num_batches = 0
@@ -146,11 +182,21 @@ def train(args):
             b_ids = b_ids.to(device)
             b_mask = b_mask.to(device)
 
-            optimizer_cse.zero_grad()
-            logits = model.bert.embed(b_ids)
-            logits = model.bert.encode(logits, b_mask)
+            # Get different embeddings with different dropout masks
+            logits_1 = model.bert(b_ids, b_mask)['pooler_output']
+            logits_2 = model.bert(b_ids, b_mask)['pooler_output']
+
+            # Calculate mean SimCSE loss function
+            loss = contrastive_loss(logits_1, logits_2)
 
+            loss.backward()
+            optimizer_cse.step()
 
+            train_loss += loss.item()
+            num_batches += 0
+
+        train_loss = train_loss / num_batches
+        print(f"Epoch {epoch}: train loss :: {train_loss :.3f}")
 
 
 def get_args():
@@ -177,18 +223,18 @@ if __name__ == "__main__":
     print('Finetuning minBERT with Unsupervised SimCSE...')
     config = SimpleNamespace(
         filepath='contrastive-nli.pt',
-        lr=args.lr,
+        lr_cse=args.lr_cse,
+        lr_classifier=args.lr_classifier,
         num_cpu_cores=args.num_cpu_cores,
         use_gpu=args.use_gpu,
         epochs=args.epochs,
         batch_size_cse=args.batch_size_cse,
         batch_size_classifier=args.batch_size_classifier,
+        hidden_dropout_prob=args.hidden_dropout_prob,
         train_bert='data/twitter-unsup.csv',
         train='data/ids-sst-train.csv',
         dev='data/ids-sst-dev.csv',
-        test='data/ids-sst-test-student.csv',
-        dev_out = 'predictions/' + args.fine_tune_mode + '-sst-dev-out.csv',
-        test_out = 'predictions/' + args.fine_tune_mode + '-sst-test-out.csv'
+        test='data/ids-sst-test-student.csv'
     )
     
     train(config)