refactor(collator): combine collator for pair and triplet

Author: LongxingTan

README.md

@@ -195,15 +195,14 @@ print(response)
 import torch.nn as nn
 from datasets import load_dataset
 from transformers import AutoTokenizer, AdamW, get_linear_schedule_with_warmup, TrainingArguments
-from retrievals import AutoModelForEmbedding, RetrievalTrainer, PairCollator, TripletCollator
+from retrievals import AutoModelForEmbedding, RetrievalTrainer, RetrievalCollator
 from retrievals.losses import ArcFaceAdaptiveMarginLoss, InfoNCE, SimCSE, TripletLoss
 
 model_name_or_path: str = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
 batch_size: int = 32
 epochs: int = 3
 
 train_dataset = load_dataset('shibing624/nli_zh', 'STS-B')['train']
-train_dataset = train_dataset.rename_columns({'sentence1': 'query', 'sentence2': 'positive'})
 tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)
 model = AutoModelForEmbedding.from_pretrained(model_name_or_path, pooling_method="mean")
 model = model.set_train_type('pairwise')
@@ -225,7 +224,7 @@ trainer = RetrievalTrainer(
     model=model,
     args=training_arguments,
     train_dataset=train_dataset,
-    data_collator=PairCollator(tokenizer, query_max_length=32, document_max_length=128),
+    data_collator=RetrievalCollator(tokenizer, keys=['sentence1', 'sentence2'], max_lengths=[32, 128]),
     loss_fn=InfoNCE(nn.CrossEntropyLoss(label_smoothing=0.05)),
 )
 trainer.optimizer = optimizer
@@ -240,7 +239,7 @@ trainer.train()
 import torch.nn as nn
 from datasets import load_dataset
 from transformers import AutoTokenizer, AdamW, get_linear_schedule_with_warmup, TrainingArguments
-from retrievals import AutoModelForEmbedding, RetrievalTrainer, PairCollator, TripletCollator
+from retrievals import AutoModelForEmbedding, RetrievalTrainer, RetrievalCollator
 from retrievals.losses import InfoNCE, SimCSE, TripletLoss
 
 def add_instructions(example):
@@ -255,7 +254,6 @@ query_instruction = "Retrieve relevant passages that answer the query\nQuery: "
 document_instruction = "Document: "
 
 train_dataset = load_dataset('shibing624/nli_zh', 'STS-B')['train']
-train_dataset = train_dataset.rename_columns({'sentence1': 'query', 'sentence2': 'positive'})
 train_dataset = train_dataset.map(add_instructions)
 tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)
 model = AutoModelForEmbedding.from_pretrained(model_name_or_path, pooling_method="last", use_lora=True)
@@ -275,7 +273,7 @@ trainer = RetrievalTrainer(
     model=model,
     args=training_arguments,
     train_dataset=train_dataset,
-    data_collator=PairCollator(tokenizer, query_max_length=64, document_max_length=128),
+    data_collator=RetrievalCollator(tokenizer, keys=['sentence1', 'sentence2'], max_lengths=[32, 128]),
 )
 trainer.optimizer = optimizer
 trainer.scheduler = scheduler

README_ja-JP.md

@@ -190,34 +190,36 @@ print(response)
 import torch.nn as nn
 from datasets import load_dataset
 from transformers import AutoTokenizer, AdamW, get_linear_schedule_with_warmup, TrainingArguments
-from retrievals import AutoModelForEmbedding, RetrievalTrainer, PairCollator, TripletCollator
+from retrievals import AutoModelForEmbedding, RetrievalTrainer, RetrievalCollator
 from retrievals.losses import ArcFaceAdaptiveMarginLoss, InfoNCE, SimCSE, TripletLoss
 
 model_name_or_path: str = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
-batch_size: int = 128
+batch_size: int = 32
 epochs: int = 3
 
 train_dataset = load_dataset('shibing624/nli_zh', 'STS-B')['train']
-train_dataset = train_dataset.rename_columns({'sentence1': 'query', 'sentence2': 'document'})
 tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)
 model = AutoModelForEmbedding.from_pretrained(model_name_or_path, pooling_method="mean")
 model = model.set_train_type('pairwise')
 
 optimizer = AdamW(model.parameters(), lr=5e-5)
 num_train_steps = int(len(train_dataset) / batch_size * epochs)
-scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0.05 * num_train_steps, num_training_steps=num_train_steps)
+scheduler = get_linear_schedule_with_warmup(
+    optimizer, num_warmup_steps=0.05 * num_train_steps, num_training_steps=num_train_steps
+)
 
 training_arguments = TrainingArguments(
     output_dir='./checkpoints',
     num_train_epochs=epochs,
     per_device_train_batch_size=batch_size,
     remove_unused_columns=False,
+    logging_steps=100,
 )
 trainer = RetrievalTrainer(
     model=model,
     args=training_arguments,
     train_dataset=train_dataset,
-    data_collator=PairCollator(tokenizer, query_max_length=128, document_max_length=128),
+    data_collator=RetrievalCollator(tokenizer, keys=['sentence1', 'sentence2'], max_lengths=[32, 128]),
     loss_fn=InfoNCE(nn.CrossEntropyLoss(label_smoothing=0.05)),
 )
 trainer.optimizer = optimizer

README_zh-CN.md

@@ -198,20 +198,17 @@ print(response)
 <details><summary> 微调向量模型 </summary>
 
 ```python
-import os
 import torch.nn as nn
 from datasets import load_dataset
 from transformers import AutoTokenizer, AdamW, get_linear_schedule_with_warmup, TrainingArguments
-from retrievals import AutoModelForEmbedding, RetrievalTrainer, PairCollator, TripletCollator
+from retrievals import AutoModelForEmbedding, RetrievalTrainer, RetrievalCollator
 from retrievals.losses import ArcFaceAdaptiveMarginLoss, InfoNCE, SimCSE, TripletLoss
-os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
 
 model_name_or_path: str = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
 batch_size: int = 32
 epochs: int = 3
 
 train_dataset = load_dataset('shibing624/nli_zh', 'STS-B')['train']
-train_dataset = train_dataset.rename_columns({'sentence1': 'query', 'sentence2': 'positive'})
 tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)
 model = AutoModelForEmbedding.from_pretrained(model_name_or_path, pooling_method="mean")
 model = model.set_train_type('pairwise')
@@ -233,7 +230,7 @@ trainer = RetrievalTrainer(
     model=model,
     args=training_arguments,
     train_dataset=train_dataset,
-    data_collator=PairCollator(tokenizer, query_max_length=32, document_max_length=128),
+    data_collator=RetrievalCollator(tokenizer, keys=['sentence1', 'sentence2'], max_lengths=[32, 128]),
     loss_fn=InfoNCE(nn.CrossEntropyLoss(label_smoothing=0.05)),
 )
 trainer.optimizer = optimizer
@@ -249,7 +246,7 @@ import os
 import torch.nn as nn
 from datasets import load_dataset
 from transformers import AutoTokenizer, AdamW, get_linear_schedule_with_warmup, TrainingArguments
-from retrievals import AutoModelForEmbedding, RetrievalTrainer, PairCollator, TripletCollator
+from retrievals import AutoModelForEmbedding, RetrievalTrainer, RetrievalCollator
 from retrievals.losses import InfoNCE, SimCSE, TripletLoss
 os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
 
@@ -285,7 +282,7 @@ trainer = RetrievalTrainer(
     model=model,
     args=training_arguments,
     train_dataset=train_dataset,
-    data_collator=PairCollator(tokenizer, query_max_length=64, document_max_length=128),
+    data_collator=RetrievalCollator(tokenizer, keys=['query', 'positive'], max_lengths=[64, 128]),
 )
 trainer.optimizer = optimizer
 trainer.scheduler = scheduler

docs/source/embed.rst

@@ -88,15 +88,14 @@ If the positive and negative examples have some noise in label, the directly poi
     import torch.nn as nn
     from datasets import load_dataset
     from transformers import AutoTokenizer, AdamW, get_linear_schedule_with_warmup, TrainingArguments
-    from retrievals import AutoModelForEmbedding, RetrievalTrainer, PairCollator, TripletCollator
+    from retrievals import AutoModelForEmbedding, RetrievalTrainer, RetrievalCollator
     from retrievals.losses import ArcFaceAdaptiveMarginLoss, InfoNCE, SimCSE, TripletLoss
 
     model_name_or_path: str = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
     batch_size: int = 128
     epochs: int = 3
 
     train_dataset = load_dataset('shibing624/nli_zh', 'STS-B')['train']
-    train_dataset = train_dataset.rename_columns({'sentence1': 'query', 'sentence2': 'positive'})
     tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)
     model = AutoModelForEmbedding.from_pretrained(model_name_or_path, pooling_method="mean")
     optimizer = AdamW(model.parameters(), lr=5e-5)
@@ -113,7 +112,7 @@ If the positive and negative examples have some noise in label, the directly poi
         model=model,
         args=training_arguments,
         train_dataset=train_dataset,
-        data_collator=PairCollator(tokenizer, query_max_length=128, document_max_length=128),
+        data_collator=RetrievalCollator(tokenizer, keys=['sentence1', 'sentence2'], max_lengths=[64, 128]),,
         loss_fn=InfoNCE(nn.CrossEntropyLoss(label_smoothing=0.05)),
     )
     trainer.optimizer = optimizer
@@ -136,6 +135,7 @@ If the positive and negative examples have some noise in label, the directly poi
       --model_name_or_path $MODEL_NAME \
       --do_train \
       --data_name_or_path $TRAIN_DATA \
+      --query_key query \
       --positive_key positive \
       --negative_key negative \
       --learning_rate 3e-5 \
@@ -167,6 +167,7 @@ If the positive and negative examples have some noise in label, the directly poi
       --pooling_method last \
       --do_train \
       --data_name_or_path $TRAIN_DATA \
+      --query_key query \
       --positive_key positive \
       --negative_key negative \
       --use_lora True \

docs/source/quick-start.rst

@@ -52,7 +52,7 @@ To further improve the retrieval performance, we can fine tune the embedding mod
     import torch.nn as nn
     from datasets import load_dataset
     from transformers import AutoTokenizer, AdamW, get_linear_schedule_with_warmup, TrainingArguments
-    from retrievals import AutoModelForEmbedding, RetrievalTrainer, PairCollator, TripletCollator
+    from retrievals import AutoModelForEmbedding, RetrievalTrainer, RetrievalCollator, RetrievalCollator
     from retrievals.losses import ArcFaceAdaptiveMarginLoss, InfoNCE, SimCSE, TripletLoss
 
     model_name_or_path: str = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
@@ -77,7 +77,7 @@ To further improve the retrieval performance, we can fine tune the embedding mod
         model=model,
         args=training_arguments,
         train_dataset=train_dataset,
-        data_collator=PairCollator(tokenizer, query_max_length=128, document_max_length=128),
+        data_collator=RetrievalCollator(tokenizer, keys=['query', 'document'], max_lengths=[128, 128]),
         loss_fn=InfoNCE(nn.CrossEntropyLoss(label_smoothing=0.05)),
     )
     trainer.optimizer = optimizer

examples/0_embedding/README.md

@@ -36,6 +36,7 @@ from Cython.Compiler.Options import embed
 
 Train directly using shell script, refer to the [document](https://open-retrievals.readthedocs.io/en/master/embed.html)
 
+
 ### Transformer encoder embedding
 
 Refer to [the fine-tuning code](./train_pairwise.py) to train the model like

examples/0_embedding/train_llm.py

@@ -22,8 +22,8 @@
     AutoModelForEmbedding,
     AutoModelForRetrieval,
     PairwiseModel,
+    RetrievalCollator,
     RetrievalTrainer,
-    TripletCollator,
 )
 from retrievals.losses import InfoNCE, TripletLoss, TripletRankingLoss
 
@@ -259,8 +259,10 @@ def main():
         model=model,
         args=training_args,
         train_dataset=train_dataset,
-        data_collator=TripletCollator(
-            tokenizer, query_max_length=data_args.query_max_length, document_max_length=data_args.document_max_length
+        data_collator=RetrievalCollator(
+            tokenizer,
+            keys=['query', 'positive', 'negative'],
+            max_lengths=[data_args.query_max_length, data_args.document_max_length, data_args.document_max_length],
         ),
     )
     trainer.optimizer = optimizer

examples/0_embedding/train_pairwise.py

@@ -9,7 +9,7 @@
     get_linear_schedule_with_warmup,
 )
 
-from retrievals import AutoModelForEmbedding, PairCollator, RetrievalTrainer
+from retrievals import AutoModelForEmbedding, RetrievalCollator, RetrievalTrainer
 from retrievals.losses import InfoNCE, SimCSE, TripletLoss
 
 model_name_or_path: str = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
@@ -41,7 +41,7 @@ def train():
         model=model,
         args=training_args,
         train_dataset=train_dataset,
-        data_collator=PairCollator(tokenizer, query_max_length=128, document_max_length=128),
+        data_collator=RetrievalCollator(tokenizer, keys=['query', 'positive'], max_lengths=[128, 128]),
         loss_fn=InfoNCE(nn.CrossEntropyLoss(label_smoothing=0.05)),
     )
     trainer.optimizer = optimizer

examples/eval/eval_retrieval2.py

@@ -10,9 +10,8 @@
     AutoModelForRanking,
     ColBERT,
     LLMRanker,
-    PairCollator,
+    RetrievalCollator,
     RetrievalTrainer,
-    TripletCollator,
 )
 
 logger = logging.getLogger(__name__)

src/retrievals/__init__.py

@@ -2,9 +2,8 @@
     ColBertCollator,
     EncodeCollator,
     LLMRerankCollator,
-    PairCollator,
     RerankCollator,
-    TripletCollator,
+    RetrievalCollator,
 )
 from .data.dataset import EncodeDataset, RerankTrainDataset, RetrievalTrainDataset
 from .models.embedding_auto import AutoModelForEmbedding, ListwiseModel, PairwiseModel

src/retrievals/data/__init__.py

@@ -2,9 +2,8 @@
     ColBertCollator,
     EncodeCollator,
     LLMRerankCollator,
-    PairCollator,
     RerankCollator,
-    TripletCollator,
+    RetrievalCollator,
 )
 from .dataset import EncodeDataset, RerankTrainDataset, RetrievalTrainDataset
 from .sampler import GroupedBatchSampler, GroupSortedBatchSampler