refactor(embed): construct the embed model

Author: LongxingTan

README.md

@@ -195,7 +195,7 @@ 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, RetrievalCollator
+from retrievals import AutoModelForEmbedding, RetrievalTrainer, RetrievalCollator, PairwiseModel
 from retrievals.losses import ArcFaceAdaptiveMarginLoss, InfoNCE, SimCSE, TripletLoss
 
 model_name_or_path: str = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
@@ -205,9 +205,9 @@ epochs: int = 3
 train_dataset = load_dataset('shibing624/nli_zh', 'STS-B')['train']
 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')
+train_model = PairwiseModel(model)
 
-optimizer = AdamW(model.parameters(), lr=5e-5)
+optimizer = AdamW(train_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
@@ -221,7 +221,7 @@ training_arguments = TrainingArguments(
     logging_steps=100,
 )
 trainer = RetrievalTrainer(
-    model=model,
+    model=train_model,
     args=training_arguments,
     train_dataset=train_dataset,
     data_collator=RetrievalCollator(tokenizer, keys=['sentence1', 'sentence2'], max_lengths=[32, 128]),
@@ -239,26 +239,26 @@ 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, RetrievalCollator
+from retrievals import AutoModelForEmbedding, RetrievalTrainer, RetrievalCollator, PairwiseModel
 from retrievals.losses import InfoNCE, SimCSE, TripletLoss
 
 def add_instructions(example):
-    example['query'] = query_instruction + example['query']
-    example['positive'] = document_instruction + example['positive']
+    example['query'] = query_instruction.format(example['query'])
+    example['positive'] = document_instruction.format(example['positive'])
     return example
 
 model_name_or_path: str = "Qwen/Qwen2-1.5B-Instruct"
 batch_size: int = 8
 epochs: int = 3
-query_instruction = "Retrieve relevant passages that answer the query\nQuery: "
-document_instruction = "Document: "
+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.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)
-model = model.set_train_type('pairwise', loss_fn=InfoNCE(nn.CrossEntropyLoss(label_smoothing=0.05)))
-optimizer = AdamW(model.parameters(), lr=5e-5)
+train_model = PairwiseModel(model, loss_fn=InfoNCE(nn.CrossEntropyLoss(label_smoothing=0.05)))
+optimizer = AdamW(train_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)
 
@@ -270,7 +270,7 @@ training_arguments = TrainingArguments(
     logging_steps=100,
 )
 trainer = RetrievalTrainer(
-    model=model,
+    model=train_model,
     args=training_arguments,
     train_dataset=train_dataset,
     data_collator=RetrievalCollator(tokenizer, keys=['sentence1', 'sentence2'], max_lengths=[32, 128]),
@@ -423,8 +423,8 @@ train_dataset = RetrievalTrainDataset(
     data_name_or_path='C-MTEB/T2Reranking',
     positive_key='positive',
     negative_key='negative',
-    query_instruction='A: ',
-    document_instruction='B: ',
+    query_instruction='A: {}',
+    document_instruction='B: {}',
     dataset_split='dev',
 )
 data_collator = LLMRerankCollator(tokenizer=tokenizer, max_length=max_length, prompt=task_prompt, add_target_token='Yes')

README_ja-JP.md

@@ -190,7 +190,7 @@ 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, RetrievalCollator
+from retrievals import AutoModelForEmbedding, RetrievalTrainer, RetrievalCollator, PairwiseModel
 from retrievals.losses import ArcFaceAdaptiveMarginLoss, InfoNCE, SimCSE, TripletLoss
 
 model_name_or_path: str = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
@@ -200,9 +200,9 @@ epochs: int = 3
 train_dataset = load_dataset('shibing624/nli_zh', 'STS-B')['train']
 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')
+train_model = PairwiseModel(model)
 
-optimizer = AdamW(model.parameters(), lr=5e-5)
+optimizer = AdamW(train_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
@@ -216,7 +216,7 @@ training_arguments = TrainingArguments(
     logging_steps=100,
 )
 trainer = RetrievalTrainer(
-    model=model,
+    model=train_model,
     args=training_arguments,
     train_dataset=train_dataset,
     data_collator=RetrievalCollator(tokenizer, keys=['sentence1', 'sentence2'], max_lengths=[32, 128]),
@@ -234,7 +234,7 @@ from retrievals import AutoModelForEmbedding
 model = AutoModelForEmbedding.from_pretrained(
     "mistralai/Mistral-7B-v0.1",
     pooling_method='last',
-    query_instruction=f'Instruct: Retrieve semantically similar text\nQuery: '
+    query_instruction=f'Instruct: Retrieve semantically similar text\nQuery: {}'
 )
 ```
 

README_zh-CN.md

@@ -201,7 +201,7 @@ 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, RetrievalCollator
+from retrievals import AutoModelForEmbedding, RetrievalTrainer, RetrievalCollator, PairwiseModel
 from retrievals.losses import ArcFaceAdaptiveMarginLoss, InfoNCE, SimCSE, TripletLoss
 
 model_name_or_path: str = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
@@ -211,9 +211,9 @@ epochs: int = 3
 train_dataset = load_dataset('shibing624/nli_zh', 'STS-B')['train']
 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')
+train_model = PairwiseModel(model)
 
-optimizer = AdamW(model.parameters(), lr=5e-5)
+optimizer = AdamW(train_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
@@ -227,7 +227,7 @@ training_arguments = TrainingArguments(
     logging_steps=100,
 )
 trainer = RetrievalTrainer(
-    model=model,
+    model=train_model,
     args=training_arguments,
     train_dataset=train_dataset,
     data_collator=RetrievalCollator(tokenizer, keys=['sentence1', 'sentence2'], max_lengths=[32, 128]),
@@ -246,28 +246,28 @@ 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, RetrievalCollator
+from retrievals import AutoModelForEmbedding, RetrievalTrainer, RetrievalCollator, PairwiseModel
 from retrievals.losses import InfoNCE, SimCSE, TripletLoss
 os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
 
 def add_instructions(example):
-    example['query'] = query_instruction + example['query']
-    example['positive'] = document_instruction + example['positive']
+    example['query'] = query_instruction.format(example['query'])
+    example['positive'] = document_instruction.format(example['positive'])
     return example
 
 model_name_or_path: str = "Qwen/Qwen2-1.5B-Instruct"
 batch_size: int = 8
 epochs: int = 3
-query_instruction = "Retrieve relevant passages that answer the query\nQuery: "
-document_instruction = "Document: "
+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)
-model = model.set_train_type('pairwise', loss_fn=InfoNCE(nn.CrossEntropyLoss(label_smoothing=0.05)))
-optimizer = AdamW(model.parameters(), lr=5e-5)
+train_model = PairwiseModel(model, loss_fn=InfoNCE(nn.CrossEntropyLoss(label_smoothing=0.05)))
+optimizer = AdamW(train_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)
 
@@ -279,7 +279,7 @@ training_arguments = TrainingArguments(
     logging_steps=100,
 )
 trainer = RetrievalTrainer(
-    model=model,
+    model=train_model,
     args=training_arguments,
     train_dataset=train_dataset,
     data_collator=RetrievalCollator(tokenizer, keys=['query', 'positive'], max_lengths=[64, 128]),
@@ -439,8 +439,8 @@ train_dataset = RetrievalTrainDataset(
     data_name_or_path='C-MTEB/T2Reranking',
     positive_key='positive',
     negative_key='negative',
-    query_instruction='A: ',
-    document_instruction='B: ',
+    query_instruction='A: {}',
+    document_instruction='B: {}',
     dataset_split='dev',
 )
 data_collator = LLMRerankCollator(

docs/source/embed.rst

@@ -88,31 +88,34 @@ 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, RetrievalCollator
+    from retrievals import AutoModelForEmbedding, RetrievalTrainer, RetrievalCollator, PairwiseModel
     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']
     tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)
     model = AutoModelForEmbedding.from_pretrained(model_name_or_path, pooling_method="mean")
+    train_model = PairwiseModel(model)
+
     optimizer = AdamW(model.parameters(), lr=5e-5)
-    num_train_steps=int(len(train_dataset) / batch_size * epochs)
+    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)
 
     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,
+        model=train_model,
         args=training_arguments,
         train_dataset=train_dataset,
-        data_collator=RetrievalCollator(tokenizer, keys=['sentence1', 'sentence2'], max_lengths=[64, 128]),,
+        data_collator=RetrievalCollator(tokenizer, keys=['sentence1', 'sentence2'], max_lengths=[64, 128]),
         loss_fn=InfoNCE(nn.CrossEntropyLoss(label_smoothing=0.05), use_inbatch_negative=True, negatives_cross_device=False),
     )
     trainer.optimizer = optimizer
@@ -172,8 +175,8 @@ If the positive and negative examples have some noise in label, the directly poi
       --positive_key positive \
       --negative_key negative \
       --use_lora True \
-      --query_instruction "Retrieve the possible answer for query.\nQuery: " \
-      --document_instruction 'Document: ' \
+      --query_instruction "Retrieve the possible answer for query.\nQuery: {}" \
+      --document_instruction 'Document: {}' \
       --learning_rate 2e-4 \
       --bf16 \
       --num_train_epochs 3 \

examples/0_embedding/train_llm.py

@@ -244,15 +244,16 @@ def main():
         pooling_method=training_args.pooling_method,
         lora_config=lora_config,
     )
-    model = model.set_train_type(
-        "pairwise",
+
+    train_model = PairwiseModel(
+        model,
         loss_fn=TripletLoss(
             use_inbatch_negative=training_args.use_inbatch_neg,
             negatives_cross_device=training_args.negatives_cross_device,
         ),
     )
 
-    optimizer = get_optimizer(model, lr=5e-5, weight_decay=1e-3)
+    optimizer = get_optimizer(train_model, lr=5e-5, weight_decay=1e-3)
 
     lr_scheduler = get_scheduler(
         optimizer,
@@ -262,7 +263,7 @@ def main():
     )
 
     trainer = RetrievalTrainer(
-        model=model,
+        model=train_model,
         args=training_args,
         train_dataset=train_dataset,
         data_collator=RetrievalCollator(

examples/0_embedding/train_pairwise.py

@@ -9,7 +9,12 @@
     get_linear_schedule_with_warmup,
 )
 
-from retrievals import AutoModelForEmbedding, RetrievalCollator, RetrievalTrainer
+from retrievals import (
+    AutoModelForEmbedding,
+    PairwiseModel,
+    RetrievalCollator,
+    RetrievalTrainer,
+)
 from retrievals.losses import InfoNCE, SimCSE, TripletLoss
 
 model_name_or_path: str = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
@@ -23,9 +28,9 @@ def 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')
+    train_model = PairwiseModel(model)
 
-    optimizer = AdamW(model.parameters(), lr=5e-5)
+    optimizer = AdamW(train_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
@@ -38,7 +43,7 @@ def train():
         remove_unused_columns=False,
     )
     trainer = RetrievalTrainer(
-        model=model,
+        model=train_model,
         args=training_args,
         train_dataset=train_dataset,
         data_collator=RetrievalCollator(tokenizer, keys=['query', 'positive'], max_lengths=[128, 128]),
@@ -48,7 +53,7 @@ def train():
     trainer.scheduler = scheduler
     trainer.train()
 
-    model.save_pretrained(training_args.output_dir)
+    train_model.save_pretrained(training_args.output_dir)
     if trainer.is_world_process_zero():
         tokenizer.save_pretrained(training_args.output_dir)
 

examples/README_zh_CN.md

@@ -30,7 +30,7 @@ export HF_ENDPOINT=https://hf-mirror.com
 - [向量模型pairwise微调](./0_embedding/train_pairwise.py)
 - [decoder大模型向量模型pairwise微调](./0_embedding/train_llm.py)
   - 设置 `query_instruction`
-    - "给定一个查询和一个相关文档，检索与查询相关的文档\n查询: "
+    - "给定一个查询和一个相关文档，检索与查询相关的文档\n查询: {}"
   - 使用适当的 `pooling_method`
     - `last`
   - 由于模型尺寸较大，可能需要减少批次大小

src/retrievals/data/collator.py

@@ -223,7 +223,6 @@ def __init__(
         add_target_token: str = '',
         sep_token: str = "\n",
         max_length: int = 128,
-        tokenize_args: Optional[Dict] = None,
         pad_to_multiple_of: Optional[int] = 8,
     ):
         self.tokenizer = tokenizer