chore: import upstream snapshot with attribution

2026-07-13 13:39:21 +08:00
commit bc34f6df14
1149 changed files with 328099 additions and 0 deletions
@@ -0,0 +1,87 @@
+import logging
+from typing import Tuple
+from transformers import (
+    AutoModel, AutoConfig,
+    AutoTokenizer, PreTrainedTokenizer
+)
+
+from FlagEmbedding.abc.finetune.embedder import AbsEmbedderRunner, AbsEmbedderModel, EmbedderTrainerCallbackForDataRefresh
+from .modeling import BiEncoderOnlyEmbedderModel
+from .trainer import EncoderOnlyEmbedderTrainer
+
+logger = logging.getLogger(__name__)
+
+
+class EncoderOnlyEmbedderRunner(AbsEmbedderRunner):
+    """
+    Finetune Runner for base embedding models.
+    """
+    def load_tokenizer_and_model(self) -> Tuple[PreTrainedTokenizer, AbsEmbedderModel]:
+        """Load tokenizer and model.
+
+        Returns:
+            Tuple[PreTrainedTokenizer, AbsEmbedderModel]: Tokenizer and model instances.
+        """
+        tokenizer = AutoTokenizer.from_pretrained(
+            self.model_args.model_name_or_path,
+            cache_dir=self.model_args.cache_dir,
+            token=self.model_args.token,
+            use_fast=self.model_args.use_fast_tokenizer,
+            trust_remote_code=self.model_args.trust_remote_code
+        )
+        base_model = AutoModel.from_pretrained(
+            self.model_args.model_name_or_path,
+            cache_dir=self.model_args.cache_dir,
+            token=self.model_args.token,
+            trust_remote_code=self.model_args.trust_remote_code
+        )
+
+        num_labels = 1
+        config = AutoConfig.from_pretrained(
+            self.model_args.config_name if self.model_args.config_name else self.model_args.model_name_or_path,
+            num_labels=num_labels,
+            cache_dir=self.model_args.cache_dir,
+            token=self.model_args.token,
+            trust_remote_code=self.model_args.trust_remote_code,
+        )
+        logger.info('Config: %s', config)
+
+        model = BiEncoderOnlyEmbedderModel(
+            base_model,
+            tokenizer=tokenizer,
+            negatives_cross_device=self.training_args.negatives_cross_device,
+            temperature=self.training_args.temperature,
+            sub_batch_size=self.training_args.sub_batch_size,
+            kd_loss_type=self.training_args.kd_loss_type,
+            use_mrl=self.training_args.use_mrl,
+            mrl_dims=self.training_args.mrl_dims,
+            sentence_pooling_method=self.training_args.sentence_pooling_method,
+            normalize_embeddings=self.training_args.normalize_embeddings
+        )
+
+        if self.training_args.gradient_checkpointing:
+            model.enable_input_require_grads()
+
+        if self.training_args.fix_position_embedding:
+            for k, v in model.named_parameters():
+                if "position_embeddings" in k:
+                    logging.info(f"Freeze the parameters for {k}")
+                    v.requires_grad = False
+        return tokenizer, model
+
+    def load_trainer(self) -> EncoderOnlyEmbedderTrainer:
+        """Load the trainer.
+
+        Returns:
+            EncoderOnlyEmbedderTrainer: Loaded trainer instance.
+        """
+        trainer = EncoderOnlyEmbedderTrainer(
+            model=self.model,
+            args=self.training_args,
+            train_dataset=self.train_dataset,
+            data_collator=self.data_collator,
+            processing_class=self.tokenizer
+        )
+        if self.data_args.same_dataset_within_batch:
+            trainer.add_callback(EmbedderTrainerCallbackForDataRefresh(self.train_dataset))
+        return trainer