update test data use.

Author: shibing624

examples/training_sup_cosent_en.py

@@ -39,6 +39,26 @@ def calc_similarity_scores(args, sents1, sents2, labels):
     return spearman
 
 
+def load_en_stsb_dataset(stsb_file):
+    # Convert the dataset to a DataLoader ready for training
+    logger.info("Read STSbenchmark dataset")
+    train_samples = []
+    valid_samples = []
+    test_samples = []
+    with gzip.open(stsb_file, 'rt', encoding='utf8') as f:
+        reader = csv.DictReader(f, delimiter='\t', quoting=csv.QUOTE_NONE)
+        for row in reader:
+            score = float(row['score'])
+            if row['split'] == 'dev':
+                valid_samples.append((row['sentence1'], row['sentence2'], score))
+            elif row['split'] == 'test':
+                test_samples.append((row['sentence1'], row['sentence2'], score))
+            else:
+                score = int(score > 2.5)
+                train_samples.append((row['sentence1'], row['sentence2'], score))
+    return train_samples, valid_samples, test_samples
+
+
 def main():
     parser = argparse.ArgumentParser('CoSENT Text Matching task')
     parser.add_argument('--model_name', default='bert-base-uncased', type=str, help='name of transformers model')
@@ -56,27 +76,11 @@ def main():
     args = parser.parse_args()
     logger.info(args)
 
-    test_samples = []
+    train_samples, valid_samples, test_samples = load_en_stsb_dataset(args.stsb_file)
+
     if args.do_train:
         model = CosentModel(model_name_or_path=args.model_name, encoder_type=args.encoder_type,
                             max_seq_length=args.max_seq_length)
-
-        # Convert the dataset to a DataLoader ready for training
-        logger.info("Read STSbenchmark dataset")
-        train_samples = []
-        valid_samples = []
-        test_samples = []
-        with gzip.open(args.stsb_file, 'rt', encoding='utf8') as f:
-            reader = csv.DictReader(f, delimiter='\t', quoting=csv.QUOTE_NONE)
-            for row in reader:
-                score = float(row['score'])
-                if row['split'] == 'dev':
-                    valid_samples.append((row['sentence1'], row['sentence2'], score))
-                elif row['split'] == 'test':
-                    test_samples.append((row['sentence1'], row['sentence2'], score))
-                else:
-                    train_samples.append((row['sentence1'], score))
-                    train_samples.append((row['sentence2'], score))
         train_dataset = CosentTrainDataset(model.tokenizer, train_samples, args.max_seq_length)
         valid_dataset = CosentTestDataset(model.tokenizer, valid_samples, args.max_seq_length)
         model.train(train_dataset,
@@ -86,21 +90,20 @@ def main():
                     batch_size=args.batch_size,
                     lr=args.learning_rate)
         logger.info(f"Model saved to {args.output_dir}")
+
     if args.do_predict:
         model = CosentModel(model_name_or_path=args.output_dir, encoder_type=args.encoder_type,
                             max_seq_length=args.max_seq_length)
-        test_data = test_samples
-
         # Predict embeddings
         srcs = []
         trgs = []
         labels = []
-        for terms in test_data:
+        for terms in test_samples:
             src, trg, label = terms[0], terms[1], terms[2]
             srcs.append(src)
             trgs.append(trg)
             labels.append(label)
-        logger.debug(f'{test_data[0]}')
+        logger.debug(f'{test_samples[0]}')
         sentence_embeddings = model.encode(srcs)
         logger.debug(f"{type(sentence_embeddings)}, {sentence_embeddings.shape}, {sentence_embeddings[0].shape}")
         # Predict similarity scores

examples/training_sup_sentencebert_en.py

@@ -39,14 +39,35 @@ def calc_similarity_scores(args, sents1, sents2, labels):
     return spearman
 
 
+def load_en_stsb_dataset(stsb_file):
+    # Convert the dataset to a DataLoader ready for training
+    logger.info("Read STSbenchmark dataset")
+    train_samples = []
+    valid_samples = []
+    test_samples = []
+    with gzip.open(stsb_file, 'rt', encoding='utf8') as f:
+        reader = csv.DictReader(f, delimiter='\t', quoting=csv.QUOTE_NONE)
+        for row in reader:
+            score = float(row['score'])
+            if row['split'] == 'dev':
+                valid_samples.append((row['sentence1'], row['sentence2'], score))
+            elif row['split'] == 'test':
+                test_samples.append((row['sentence1'], row['sentence2'], score))
+            else:
+                score = int(score > 2.5)
+                train_samples.append((row['sentence1'], row['sentence2'], score))
+    return train_samples, valid_samples, test_samples
+
+
 def main():
     parser = argparse.ArgumentParser('SentenceBERT Text Matching task')
     parser.add_argument('--model_name', default='bert-base-uncased', type=str, help='name of transformers model')
     parser.add_argument('--stsb_file', default='data/English-STS-B/stsbenchmark.tsv.gz', type=str,
                         help='Train data path')
     parser.add_argument("--do_train", action="store_true", help="Whether to run training.")
     parser.add_argument("--do_predict", action="store_true", help="Whether to run predict.")
-    parser.add_argument('--output_dir', default='./outputs/STS-B-en-sentencebert', type=str, help='Model output directory')
+    parser.add_argument('--output_dir', default='./outputs/STS-B-en-sentencebert', type=str,
+                        help='Model output directory')
     parser.add_argument('--max_seq_length', default=64, type=int, help='Max sequence length')
     parser.add_argument('--num_epochs', default=10, type=int, help='Number of training epochs')
     parser.add_argument('--batch_size', default=64, type=int, help='Batch size')
@@ -56,26 +77,11 @@ def main():
     args = parser.parse_args()
     logger.info(args)
 
-    test_samples = []
+    train_samples, valid_samples, test_samples = load_en_stsb_dataset(args.stsb_file)
+
     if args.do_train:
         model = SentenceBertModel(model_name_or_path=args.model_name, encoder_type=args.encoder_type,
                                   max_seq_length=args.max_seq_length)
-
-        # Convert the dataset to a DataLoader ready for training
-        logger.info("Read STSbenchmark dataset")
-        train_samples = []
-        valid_samples = []
-        with gzip.open(args.stsb_file, 'rt', encoding='utf8') as f:
-            reader = csv.DictReader(f, delimiter='\t', quoting=csv.QUOTE_NONE)
-            for row in reader:
-                score = float(row['score'])
-                if row['split'] == 'dev':
-                    valid_samples.append((row['sentence1'], row['sentence2'], score))
-                elif row['split'] == 'test':
-                    test_samples.append((row['sentence1'], row['sentence2'], score))
-                else:
-                    score = int(score > 2.5)
-                    train_samples.append((row['sentence1'], row['sentence2'], score))
         train_dataset = SentenceBertTrainDataset(model.tokenizer, train_samples, args.max_seq_length)
         valid_dataset = SentenceBertTestDataset(model.tokenizer, valid_samples, args.max_seq_length)
         model.train(train_dataset,
@@ -85,21 +91,20 @@ def main():
                     batch_size=args.batch_size,
                     lr=args.learning_rate)
         logger.info(f"Model saved to {args.output_dir}")
+
     if args.do_predict:
         model = SentenceBertModel(model_name_or_path=args.output_dir, encoder_type=args.encoder_type,
                                   max_seq_length=args.max_seq_length)
-        test_data = test_samples
-
         # Predict embeddings
         srcs = []
         trgs = []
         labels = []
-        for terms in test_data:
+        for terms in test_samples:
             src, trg, label = terms[0], terms[1], terms[2]
             srcs.append(src)
             trgs.append(trg)
             labels.append(label)
-        logger.debug(f'{test_data[0]}')
+        logger.debug(f'{test_samples[0]}')
         sentence_embeddings = model.encode(srcs)
         logger.debug(f"{type(sentence_embeddings)}, {sentence_embeddings.shape}, {sentence_embeddings[0].shape}")
         # Predict similarity scores

examples/training_unsup_cosent_en.py

@@ -39,6 +39,26 @@ def calc_similarity_scores(args, sents1, sents2, labels):
     return spearman
 
 
+def load_en_stsb_dataset(stsb_file):
+    # Convert the dataset to a DataLoader ready for training
+    logger.info("Read STSbenchmark dataset")
+    train_samples = []
+    valid_samples = []
+    test_samples = []
+    with gzip.open(stsb_file, 'rt', encoding='utf8') as f:
+        reader = csv.DictReader(f, delimiter='\t', quoting=csv.QUOTE_NONE)
+        for row in reader:
+            score = float(row['score'])
+            if row['split'] == 'dev':
+                valid_samples.append((row['sentence1'], row['sentence2'], score))
+            elif row['split'] == 'test':
+                test_samples.append((row['sentence1'], row['sentence2'], score))
+            else:
+                score = int(score > 2.5)
+                train_samples.append((row['sentence1'], row['sentence2'], score))
+    return train_samples, valid_samples, test_samples
+
+
 def main():
     parser = argparse.ArgumentParser('CoSENT Text Matching task')
     parser.add_argument('--model_name', default='bert-base-uncased', type=str, help='name of transformers model')
@@ -58,7 +78,8 @@ def main():
     args = parser.parse_args()
     logger.info(args)
 
-    test_samples = []
+    _, valid_samples, test_samples = load_en_stsb_dataset(args.stsb_file)
+
     if args.do_train:
         model = CosentModel(model_name_or_path=args.model_name, encoder_type=args.encoder_type,
                             max_seq_length=args.max_seq_length)
@@ -82,43 +103,29 @@ def main():
                         break
 
         train_dataset = CosentTrainDataset(model.tokenizer, nli_train_samples, args.max_seq_length)
-
-        # Convert the dataset to a DataLoader ready for validation
-        logger.info("Read STSbenchmark dev and test dataset")
-        valid_samples = []
-        test_samples = []
-        with gzip.open(args.stsb_file, 'rt', encoding='utf8') as f:
-            reader = csv.DictReader(f, delimiter='\t', quoting=csv.QUOTE_NONE)
-            for row in reader:
-                score = float(row['score'])
-                if row['split'] == 'dev':
-                    valid_samples.append((row['sentence1'], row['sentence2'], score))
-                elif row['split'] == 'test':
-                    test_samples.append((row['sentence1'], row['sentence2'], score))
-
         valid_dataset = CosentTestDataset(model.tokenizer, valid_samples, args.max_seq_length)
+
         model.train(train_dataset,
                     args.output_dir,
                     eval_dataset=valid_dataset,
                     num_epochs=args.num_epochs,
                     batch_size=args.batch_size,
                     lr=args.learning_rate)
         logger.info(f"Model saved to {args.output_dir}")
+
     if args.do_predict:
         model = CosentModel(model_name_or_path=args.output_dir, encoder_type=args.encoder_type,
                             max_seq_length=args.max_seq_length)
-        test_data = test_samples
-
         # Predict embeddings
         srcs = []
         trgs = []
         labels = []
-        for terms in test_data:
+        for terms in test_samples:
             src, trg, label = terms[0], terms[1], terms[2]
             srcs.append(src)
             trgs.append(trg)
             labels.append(label)
-        logger.debug(f'{test_data[0]}')
+        logger.debug(f'{test_samples[0]}')
         sentence_embeddings = model.encode(srcs)
         logger.debug(f"{type(sentence_embeddings)}, {sentence_embeddings.shape}, {sentence_embeddings[0].shape}")
         # Predict similarity scores