angle_emb ========= .. py:module:: angle_emb Submodules ---------- .. toctree:: :maxdepth: 1 /autoapi/angle_emb/angle/index /autoapi/angle_emb/angle_trainer/index /autoapi/angle_emb/base/index /autoapi/angle_emb/evaluation/index /autoapi/angle_emb/loss/index /autoapi/angle_emb/utils/index /autoapi/angle_emb/version/index Attributes ---------- .. autoapisummary:: angle_emb.logger angle_emb.__version__ angle_emb.__version__ Classes ------- .. autoapisummary:: angle_emb.AngleBase angle_emb.CorrelationEvaluator angle_emb.AngleBase angle_emb.CorrelationEvaluator angle_emb.Prompts angle_emb.AngleDataCollator angle_emb.Pooler angle_emb.AngleTrainer angle_emb.AngleESETrainer angle_emb.AngleLoss angle_emb.AnglE angle_emb.EvaluateCallback Functions --------- .. autoapisummary:: angle_emb.categorical_crossentropy_loss angle_emb.cosine_loss angle_emb.angle_loss angle_emb.in_batch_negative_loss angle_emb.contrastive_with_negative_loss angle_emb.angle_loss angle_emb.contrastive_with_negative_loss angle_emb.cosine_loss angle_emb.in_batch_negative_loss angle_emb.find_all_linear_names angle_emb.get_pooling angle_emb.set_device angle_emb.detect_dataset_format Package Contents ---------------- .. py:function:: categorical_crossentropy_loss(y_true: torch.Tensor, y_pred: torch.Tensor, from_logits: bool = True) -> torch.Tensor Compute categorical crossentropy :param y_true: torch.Tensor, ground truth :param y_pred: torch.Tensor, model output :param from_logits: bool, `True` means y_pred has not transformed by softmax, default True :return: torch.Tensor, loss value .. py:function:: cosine_loss(y_true: torch.Tensor, y_pred: torch.Tensor, tau: float = 20.0) -> torch.Tensor Compute cosine loss :param y_true: torch.Tensor, ground truth. The y_true must be zigzag style, such as [x[0][0], x[0][1], x[1][0], x[1][1], ...], where (x[0][0], x[0][1]) stands for a pair. :param y_pred: torch.Tensor, model output. The y_pred must be zigzag style, such as [o[0][0], o[0][1], o[1][0], o[1][1], ...], where (o[0][0], o[0][1]) stands for a pair. :param tau: float, scale factor, default 20 :return: torch.Tensor, loss value .. py:function:: angle_loss(y_true: torch.Tensor, y_pred: torch.Tensor, tau: float = 1.0, pooling_strategy: str = 'sum') Compute angle loss :param y_true: torch.Tensor, ground truth. The y_true must be zigzag style, such as [x[0][0], x[0][1], x[1][0], x[1][1], ...], where (x[0][0], x[0][1]) stands for a pair. :param y_pred: torch.Tensor, model output. The y_pred must be zigzag style, such as [o[0][0], o[0][1], o[1][0], o[1][1], ...], where (o[0][0], o[0][1]) stands for a pair. :param tau: float, scale factor, default 1.0 :return: torch.Tensor, loss value .. py:function:: in_batch_negative_loss(y_true: torch.Tensor, y_pred: torch.Tensor, tau: float = 20.0, negative_weights: float = 0.0) -> torch.Tensor Compute in-batch negative loss, i.e., contrastive loss :param y_true: torch.Tensor, ground truth. The y_true must be zigzag style, such as [x[0][0], x[0][1], x[1][0], x[1][1], ...], where (x[0][0], x[0][1]) stands for a pair. :param y_pred: torch.Tensor, model output. The y_pred must be zigzag style, such as [o[0][0], o[0][1], o[1][0], o[1][1], ...], where (o[0][0], o[0][1]) stands for a pair. :param tau: float, scale factor, default 20.0 :param negative_weights: float, negative weights, default 0.0 :return: torch.Tensor, loss value .. py:function:: contrastive_with_negative_loss(text: torch.Tensor, pos: torch.Tensor, neg: Optional[torch.Tensor] = None, tau: float = 20.0) -> torch.Tensor Compute contrastive with negative loss :param text: torch.Tensor, text. :param pos: torch.Tensor, positive samples of text. :param neg: torch.Tensor, negative samples of text. :param tau: float, scale factor, default 20.0 :return: torch.Tensor, loss value .. py:class:: AngleBase .. py:method:: encode() :abstractmethod: .. py:method:: fit() :abstractmethod: .. py:class:: CorrelationEvaluator(text1: List[str], text2: List[str], labels: List[float], batch_size: int = 32) Bases: :py:obj:`object` .. py:attribute:: text1 .. py:attribute:: text2 .. py:attribute:: labels .. py:attribute:: batch_size :value: 32 .. py:method:: __call__(model: angle_emb.base.AngleBase, show_progress: bool = True, **kwargs) -> dict Evaluate the model on the given dataset. :param model: AnglE, the model to evaluate. :param show_progress: bool, whether to show a progress bar during evaluation. :param kwargs: Additional keyword arguments to pass to the `encode` method of the model. :return: dict, The evaluation results. .. py:method:: list_all_metrics() -> List[str] Get a list of all the metrics that can be computed by this evaluator. :return: List[str], A list of all the metrics that can be computed by this evaluator. .. py:class:: AngleBase .. py:method:: encode() :abstractmethod: .. py:method:: fit() :abstractmethod: .. py:class:: CorrelationEvaluator(text1: List[str], text2: List[str], labels: List[float], batch_size: int = 32) Bases: :py:obj:`object` .. py:attribute:: text1 .. py:attribute:: text2 .. py:attribute:: labels .. py:attribute:: batch_size :value: 32 .. py:method:: __call__(model: angle_emb.base.AngleBase, show_progress: bool = True, **kwargs) -> dict Evaluate the model on the given dataset. :param model: AnglE, the model to evaluate. :param show_progress: bool, whether to show a progress bar during evaluation. :param kwargs: Additional keyword arguments to pass to the `encode` method of the model. :return: dict, The evaluation results. .. py:method:: list_all_metrics() -> List[str] Get a list of all the metrics that can be computed by this evaluator. :return: List[str], A list of all the metrics that can be computed by this evaluator. .. py:function:: angle_loss(y_true: torch.Tensor, y_pred: torch.Tensor, tau: float = 1.0, pooling_strategy: str = 'sum') Compute angle loss :param y_true: torch.Tensor, ground truth. The y_true must be zigzag style, such as [x[0][0], x[0][1], x[1][0], x[1][1], ...], where (x[0][0], x[0][1]) stands for a pair. :param y_pred: torch.Tensor, model output. The y_pred must be zigzag style, such as [o[0][0], o[0][1], o[1][0], o[1][1], ...], where (o[0][0], o[0][1]) stands for a pair. :param tau: float, scale factor, default 1.0 :return: torch.Tensor, loss value .. py:function:: contrastive_with_negative_loss(text: torch.Tensor, pos: torch.Tensor, neg: Optional[torch.Tensor] = None, tau: float = 20.0) -> torch.Tensor Compute contrastive with negative loss :param text: torch.Tensor, text. :param pos: torch.Tensor, positive samples of text. :param neg: torch.Tensor, negative samples of text. :param tau: float, scale factor, default 20.0 :return: torch.Tensor, loss value .. py:function:: cosine_loss(y_true: torch.Tensor, y_pred: torch.Tensor, tau: float = 20.0) -> torch.Tensor Compute cosine loss :param y_true: torch.Tensor, ground truth. The y_true must be zigzag style, such as [x[0][0], x[0][1], x[1][0], x[1][1], ...], where (x[0][0], x[0][1]) stands for a pair. :param y_pred: torch.Tensor, model output. The y_pred must be zigzag style, such as [o[0][0], o[0][1], o[1][0], o[1][1], ...], where (o[0][0], o[0][1]) stands for a pair. :param tau: float, scale factor, default 20 :return: torch.Tensor, loss value .. py:function:: in_batch_negative_loss(y_true: torch.Tensor, y_pred: torch.Tensor, tau: float = 20.0, negative_weights: float = 0.0) -> torch.Tensor Compute in-batch negative loss, i.e., contrastive loss :param y_true: torch.Tensor, ground truth. The y_true must be zigzag style, such as [x[0][0], x[0][1], x[1][0], x[1][1], ...], where (x[0][0], x[0][1]) stands for a pair. :param y_pred: torch.Tensor, model output. The y_pred must be zigzag style, such as [o[0][0], o[0][1], o[1][0], o[1][1], ...], where (o[0][0], o[0][1]) stands for a pair. :param tau: float, scale factor, default 20.0 :param negative_weights: float, negative weights, default 0.0 :return: torch.Tensor, loss value .. py:function:: find_all_linear_names(model: transformers.PreTrainedModel, linear_type: Optional[object] = None) -> List[str] Find all linear layer names :param model: PreTrainedModel :param linear_type: Optional[object] = None, linear type, such as nn.Linear and bnb.nn.Linear4bit. :return: List[str], linear layer names .. py:function:: get_pooling(outputs: torch.Tensor, inputs: Dict, pooling_strategy: str, padding_side: str) -> torch.Tensor Pooling the model outputs. :param outputs: torch.Tensor. Model outputs (without pooling) :param inputs: Dict. Model inputs :param pooling_strategy: str. Pooling strategy [`cls`, `cls_avg`, `cls_max`, `last`, `avg`, `mean`, `max`, `all`, int] :param padding_side: str. Padding strategy of tokenizers (`left` or `right`). It can be obtained by `tokenizer.padding_side`. .. py:data:: logger .. py:function:: set_device() -> str Set device automatically :return: str, device name .. py:data:: __version__ :value: '0.6.1' .. py:function:: detect_dataset_format(ds: datasets.Dataset) -> str Detect dataset format from raw data .. py:class:: Prompts Predefined prompts. Follow the model usage to choose the corresponding prompt. Example:: from angle_emb import Prompts # list all pre-defined prompts print(Prompts.list_prompts()) # set prompt angle.encode(*, prompt=Prompts.A) .. py:attribute:: A :value: 'Summarize sentence "{text}" in one word:"' .. py:attribute:: B :value: 'You can only output one word. Summarize "{text}":"' .. py:attribute:: C :value: 'Represent this sentence for searching relevant passages: {text}' .. py:method:: list_prompts() :classmethod: .. py:class:: AngleDataCollator AngleDataCollator. It will be implicitly used in AnglE.fit(). It handles raw data, tokenizes it, and prepares batches. :param tokenizer: PreTrainedTokenizerBase :param padding: Union[bool, str, PaddingStrategy], padding strategy :param max_length: Optional[int], max length :param return_tensors: str :param filter_duplicate: bool. Whether filter duplicate data :param text_prompt: Optional[str], prompt for text1 and text2 (format A only). Default None :param query_prompt: Optional[str], prompt for query field. Default None :param doc_prompt: Optional[str], prompt for positive/negative fields. Default None :param dataset_format: Optional[str]. Specify dataset_format: 'A', 'B', or 'C'. Default None. .. py:attribute:: tokenizer :type: transformers.tokenization_utils_base.PreTrainedTokenizerBase .. py:attribute:: padding :type: Union[bool, str, transformers.utils.PaddingStrategy] :value: 'longest' .. py:attribute:: max_length :type: Optional[int] :value: None .. py:attribute:: return_tensors :type: str :value: 'pt' .. py:attribute:: filter_duplicate :type: bool :value: True .. py:attribute:: text_prompt :type: Optional[str] :value: None .. py:attribute:: query_prompt :type: Optional[str] :value: None .. py:attribute:: doc_prompt :type: Optional[str] :value: None .. py:attribute:: dataset_format :type: Optional[str] :value: None .. py:method:: sample_from_list(text: Union[str, List[str]]) -> str :staticmethod: Sample one string from list or return string as is .. py:method:: __call__(features: List[Dict], return_tensors: str = 'pt') -> Dict[str, torch.Tensor] Collate function that handles raw data. :param features: List[Dict]. Raw data samples :param return_tensors: str. Default "pt" :return: Dict[str, torch.Tensor]. Collated data .. py:class:: Pooler(model: transformers.PreTrainedModel, pooling_strategy: Optional[Union[int, str]] = None, padding_side: Optional[str] = None) Using Pooler to obtain sentence embeddings. :param model: PreTrainedModel :param pooling_strategy: Optional[str]. Currently support [`cls`, `cls_avg`, `cls_max`, `last`, `avg`, `mean`, `max`, `all`, int]. Default None. :param padding_side: Optional[str]. `left` or `right`. Default None. :param is_llm: bool. Default False .. py:attribute:: model .. py:attribute:: pooling_strategy :value: None .. py:attribute:: padding_side :value: None .. py:method:: __call__(inputs: Dict, layer_index: int = -1, embedding_start: Optional[int] = None, embedding_size: Optional[int] = None, return_all_layer_outputs: bool = False, pooling_strategy: Optional[Union[int, str]] = None, return_mlm_logits: bool = False) -> torch.Tensor Get sentence embeddings. :param inputs: Dict. Model inputs. :param layer_index: Optional[int]. Get embeddings from specific layer. :param embedding_start: Optional[int]. Start index of embeddings. :param embedding_size: int. Set embedding size for sentence embeddings. :param return_all_layer_outputs: bool. Return all layer outputs or not. Default False. :param pooling_strategy: Optional[str]. Currently support [`cls`, `cls_avg`, `cls_max`, `last`, `avg`, `mean`, `max`, `all`, int]. Default None. :param return_mlm_logits: bool. Return logits or not. Default False. .. py:class:: AngleTrainer(pooler: Pooler, loss_kwargs: Optional[Dict] = None, dataset_format: str = 'A', teacher_name_or_path: Optional[str] = None, teacher_pooling_strategy: str = 'cls', pad_token_id: int = 0, model_kwargs: Optional[Dict] = None, **kwargs) Bases: :py:obj:`transformers.Trainer` Custom Huggingface Trainer for AnglE. :param pooler: Pooler. Required :param loss_kwargs: Optional[Dict]. Default None. :param dataset_format: str. Default 'A' :param teacher_name_or_path: Optional[str]. For distribution alignment. :param **kwargs: other parameters of Trainer. .. py:attribute:: pooler .. py:attribute:: pad_token_id :value: 0 .. py:attribute:: model_kwargs :value: None .. py:attribute:: loss_fct .. py:attribute:: teacher_name_or_path :value: None .. py:attribute:: teacher_pooling_strategy :value: 'cls' .. py:method:: compute_distillation_loss(inputs: torch.Tensor, targets: torch.Tensor, mse_weight: float = 1.0, kl_temperature: float = 1.0) -> torch.Tensor Compute distillation loss. :param inputs: torch.Tensor. Input tensor. :param targets: torch.Tensor. Target tensor. :param mse_weight: float. MSE weight. Default 1.0. :param kl_temperature: float. KL temperature. Default 1.0. :return: torch.Tensor. Distillation loss. .. py:method:: compute_mlm_loss(logits, mask_target_labels) .. py:method:: compute_loss(model, inputs, return_outputs: bool = False, **kwargs) Compute loss for AnglE. :param model: Huggingface model. :param inputs: Dict. Model inputs. :param return_outputs: bool. Return outputs or not. Default False. :return: torch.Tensor. Loss. .. py:method:: prediction_step(model, inputs, *args, **kwargs) .. py:class:: AngleESETrainer(pooler: Pooler, loss_kwargs: Optional[Dict] = None, dataset_format: str = 'A', teacher_name_or_path: Optional[str] = None, ese_kl_temperature: float = 1.0, ese_compression_size: int = 128, apply_ese_pca: bool = True, **kwargs) Bases: :py:obj:`AngleTrainer` Custom Huggingface Trainer for AnglE Espresso. :param pooler: Pooler. Required :param loss_kwargs: Optional[Dict]. Default None. :param dataset_format: str. Default 'A' :param teacher_name_or_path: Optional[str]. For distribution alignment. :param **kwargs: other parameters of Trainer. .. py:attribute:: ese_kl_temperature :value: 1.0 .. py:attribute:: ese_compression_size :value: 128 .. py:attribute:: apply_ese_pca :value: True .. py:attribute:: n_layers .. py:method:: pca_compress(m: torch.Tensor, k: int) -> torch.Tensor Get topk feature via PCA. :param m: torch.Tensor. Input tensor. :param k: int. Top-k feature size. :return: torch.Tensor. Top-k feature. .. py:method:: compute_student_loss(inputs: Dict, all_layer_outputs: torch.Tensor, labels: torch.Tensor, pooling_strategy: str, padding_side: str) -> torch.Tensor .. py:method:: compute_loss(model, inputs, return_outputs=False, **kwargs) Compute loss for Espresso. :param model: Huggingface model. :param inputs: Dict. Model inputs. :param return_outputs: bool. Return outputs or not. Default False. :return: torch.Tensor. Loss. .. py:class:: AngleLoss(cosine_w: float = 0.0, ibn_w: float = 1.0, cln_w: float = 1.0, angle_w: float = 0.02, cosine_tau: float = 20.0, ibn_tau: float = 20.0, angle_tau: float = 20.0, angle_pooling_strategy: str = 'sum', dataset_format: Optional[str] = None, **kwargs) Configure AngleLoss. :param cosine_w: float. weight for cosine_loss. Default 1.0 :param ibn_w: float. weight for in batch negative loss. Default 1.0 :param cln_w: float. weight for contrastive learning with hard negative. Default 1.0 :param angle_w: float. weight for angle loss. Default 1.0 :param cosine_tau: float. tau for cosine loss. Default 20.0 :param ibn_tau: float. tau for in batch negative loss. Default 20.0 :param angle_tau: float. tau for angle loss. Default 20.0 :param angle_pooling_strategy: str. pooling strategy for angle loss. Default'sum'. :param dataset_format: Optional[str]. Default None. .. py:attribute:: cosine_w :value: 0.0 .. py:attribute:: ibn_w :value: 1.0 .. py:attribute:: cln_w :value: 1.0 .. py:attribute:: angle_w :value: 0.02 .. py:attribute:: cosine_tau :value: 20.0 .. py:attribute:: ibn_tau :value: 20.0 .. py:attribute:: angle_tau :value: 20.0 .. py:attribute:: angle_pooling_strategy :value: 'sum' .. py:attribute:: dataset_format :value: None .. py:method:: __call__(labels: torch.Tensor, outputs: torch.Tensor) -> torch.Tensor Compute loss for AnglE. :param labels: torch.Tensor. Labels. :param outputs: torch.Tensor. Outputs. :return: torch.Tensor. Loss. .. py:class:: AnglE(model_name_or_path: str, tokenizer_name_or_path: Optional[str] = None, max_length: int = 512, model_kwargs: Optional[Dict] = None, lora_config_kwargs: Optional[Dict] = None, pooling_strategy: Optional[str] = None, apply_lora: Optional[bool] = None, train_mode: bool = True, load_kbit: Optional[int] = None, is_llm: Optional[bool] = None, pretrained_model_path: Optional[str] = None, pretrained_lora_path: Optional[str] = None, torch_dtype: Optional[torch.dtype] = None, device: Optional[str] = None, kbit_kwargs: Optional[Dict] = None, tokenizer_padding_side: Optional[str] = None, apply_billm: bool = False, billm_model_class: Optional[str] = None, load_mlm_model: bool = False, **kwargs: Any) Bases: :py:obj:`angle_emb.base.AngleBase` AnglE. Everything is heređź‘‹ :param model_name_or_path: str, model name or path. :param tokenizer_name_or_path: Optional[str]. Default None. When it set to None, it will use the same as `model_name_or_path`. :param max_length: int. Default 512 :param model_kwargs: Optional[Dict]. kwargs for model. :param lora_config_kwargs: Optional[Dict]. kwargs for peft lora_config. details refer to: https://huggingface.co/docs/peft/tutorial/peft_model_config :param pooling_strategy: Optional[str]. Pooling strategy. Currently support [`cls`, `cls_avg`, `cls_max`, `last`, `avg`, `mean`, `max`, `all`, int] :param apply_lora: Optional[bool]. Whether apply lora. Default None. :param train_mode: bool. Whether load for training. Default True. :param load_kbit: Optional[int]. Specify kbit training from [4, 8, 16]. Default None. :param is_llm: Optional[bool]. Whether the model is llm. Must be set manually. Default None. :param pretrained_model_path: Optional[str]. Default None. :param pretrained_lora_path: Optional[str]. Default None. :param torch_dtype: Optional[torch.dtype]. Specify torch_dtype. Default None. :param device: Optional[str]. Specify device. Default None. :param kbit_kwargs: Optional[Dict]. kwargs for kbit. Default None. details refer to: https://huggingface.co/docs/peft/package_reference/peft_model#peft.prepare_model_for_kbit_training :param tokenizer_padding_side: Optional[str]. Specify tokenizer padding side from [`left`, `right`]. Default None. :param apply_billm: bool. Whether apply billm. Default False. :param billm_model_class: Optional[str]. Specify billm model class. Default None. :param load_mlm_model: bool. Whether load mlm model. Default False. If set True, it will load model with AutoModelForMaskedLM. :param **kwargs: Any. .. py:attribute:: cfg_file_name :value: 'angle_config.json' .. py:attribute:: special_columns :value: ['labels'] .. py:attribute:: max_length :value: 512 .. py:attribute:: train_mode :value: True .. py:attribute:: pooling_strategy :value: None .. py:attribute:: load_kbit :value: None .. py:attribute:: is_llm :value: None .. py:attribute:: load_mlm_model :value: False .. py:attribute:: apply_lora :value: None .. py:attribute:: model_kwargs .. py:attribute:: tokenizer .. py:attribute:: pooler .. py:attribute:: __cfg .. py:method:: cuda() .. py:method:: to(device: Any) .. py:method:: kbit_post_handle(model: torch.nn.Module) -> torch.nn.Module :staticmethod: .. py:method:: find_pth_path(dirpath: str, config: Dict) -> str :staticmethod: .. py:method:: from_pretrained(model_name_or_path: str, pretrained_model_path: Optional[str] = None, pretrained_lora_path: Optional[str] = None, is_llm: Optional[bool] = None, pooling_strategy: str = 'cls', train_mode: bool = False, model_kwargs: Optional[Dict] = None, **kwargs) :staticmethod: Load AnglE from pretrained model. :param model_name_or_path: str, model name or path. Required. :param pretrained_model_path: Optional[str]. :param pretrained_lora_path: Optional[str]. :param is_llm: Optional[bool]. :param pooling_strategy: str. Pooling Strategy. Default `cls`. :param train_mode: bool. Default False. :param kwargs: Other kwargs for AnglE. :return: AnglE object. Example:: from angle_emb import AnglE angle = AnglE.from_pretrained(model_name_or_path) # fit angle.fit(*args, **kwargs) # inference angle.encode(*args, **kwargs) .. py:method:: load_config(fpath: str) -> Dict :staticmethod: .. py:method:: save_config(fpath: str) .. py:method:: fit(train_ds: datasets.Dataset, valid_ds: Optional[datasets.Dataset] = None, valid_ds_for_callback: Optional[datasets.Dataset] = None, batch_size: int = 32, output_dir: Optional[str] = None, epochs: int = 1, learning_rate: float = 1e-05, warmup_steps: int = 1000, logging_steps: int = 10, eval_steps: int = 1000, eval_strategy: str = 'steps', save_steps: int = 100, save_strategy: str = 'steps', save_total_limit: int = 1, gradient_accumulation_steps: int = 1, fp16: Optional[bool] = None, bf16: Optional[bool] = None, argument_kwargs: Optional[Dict] = None, trainer_kwargs: Optional[Dict] = None, loss_kwargs: Optional[Dict] = None, apply_ese: bool = False, filter_duplicate: bool = True, push_to_hub: bool = False, hub_model_id: Optional[str] = None, hub_private_repo: bool = True, padding: str = 'longest', text_prompt: Optional[str] = None, query_prompt: Optional[str] = None, doc_prompt: Optional[str] = None) Fit using AnglE. :param train_ds: Dataset. Raw train dataset (not tokenized). Required. :param valid_ds: Optional[Dataset]. Raw valid dataset (not tokenized). Default None. :param valid_ds_for_callback: Optional[Dataset]. Raw valid dataset for callback use (not tokenized). The dataset format should be format A. The spearmans' correlation will be computed after each epoch training and the best model will be saved. Default None. :param batch_size: int. Default 32. :param output_dir: Optional[str]. save dir. Default None. :param epochs: int. Default 1. :param learning_rate: float. Default 1e-5. :param warmup_steps: int. Default 1000. :param logging_steps: int. Default 10. :param eval_steps: int. Default 1000. :param eval_strategy: str. Default 'steps'. :param save_steps: int. Default 100. :param save_strategy: str. Default steps. :param save_total_limit: int. Default 10. :param gradient_accumulation_steps: int. Default 1. :param fp16: Optional[bool]. Default None. :param bf16: Optional[bool]. Default None. :param argument_kwargs: Optional[Dict]. kwargs for TrainingArguments. refer to: https://huggingface.co/docs/transformers/v4.37.0/en/main_classes/trainer#transformers.TrainingArguments :param trainer_kwargs: Optional[Dict]. kwargs for AngleTrainer. :param loss_kwargs: Optional[Dict]. kwargs for AngleLoss. :param apply_ese: bool, whether apply ESE training. :param filter_duplicate: bool, whether filter duplicate samples. :param push_to_hub: bool, whether push to hub. :param hub_model_id: Optional[str], hub model id. :param hub_private_repo: bool, whether push to private repo. :param padding: str, padding strategy of tokenizer. Default 'longest'. :param text_prompt: Optional[str], prompt for text1 and text2 (format A only). Default None. :param query_prompt: Optional[str], prompt template for query. Default None. :param doc_prompt: Optional[str], prompt template for documents (positive/negative). Default None. .. py:method:: evaluate(ds: datasets.Dataset, batch_size: int = 32, metric: str = 'spearman_cosine', prompt: Optional[str] = None) -> float evaluate :param data: Dataset, format A is required :param batch_size: int. Default 32. :param metric: str. Default 'spearman_cosine'. :return: float. .. py:method:: truncate_layer(layer_index: int) truncate layer :param layer_index: int. layers after layer_index will be truncated. :return: self .. py:method:: encode(inputs: Union[List[str], Tuple[str], str], max_length: Optional[int] = None, to_numpy: bool = True, embedding_start: int = 0, embedding_size: Optional[int] = None, device: Optional[Any] = None, prompt: Optional[str] = None, normalize_embedding: bool = False, padding: str = 'longest') encode texts. :param inputs: Union[List[str], Tuple[str], List[Dict], str]. Input texts. Required. :param max_length: Optional[int]. Default None. :param to_numpy: bool. Default True. :param embedding_start: int. Specify the start position of the embedding (for Espresso). :param embedding_size: Optional[int]. Specify embedding size (for Espresso). The embeddings from embedding_start to embedding_start+embedding_size will be returned. :param device: Optional[Any]. Default None. :param prompt: Optional[str]. Default None. :param normalize_embedding: bool. Default False. :param padding: str. Padding strategy of tokenizer. Default 'longest'. .. py:method:: push_to_hub(hub_model_id: str, private: bool = True, exist_ok: bool = False, **kwargs) push model to hub :param hub_model_id: str, hub model id. :param private: bool, whether push to private repo. Default True. :param exist_ok: bool, whether allow overwrite. Default False. :param kwargs: other kwargs for `push_to_hub` method. .. py:method:: save_pretrained(output_dir: str, exist_ok: bool = True) save model and tokenizer :param output_dir: str, output dir. :param exist_ok: bool, whether allow overwrite. Default True. .. py:class:: EvaluateCallback(model: AnglE, valid_ds: datasets.Dataset, evaluate_fn: Callable, save_dir: Optional[str] = None, push_to_hub: bool = False, hub_model_id: Optional[str] = None, hub_private_repo: bool = True) Bases: :py:obj:`transformers.TrainerCallback` Custom TrainerCallback for Angle. This callback will compute corrcoef for each epoch. :param model: PreTrainedModel. :param valid_ds: Dataset. :param evaluate_fn: Callable. It will receive valid_ds as input like `evaluate_fn(valid_ds)`. :param save_dir: Optional[str]. specify dir to save model with best results. .. py:attribute:: model .. py:attribute:: valid_ds .. py:attribute:: evaluate_fn .. py:attribute:: save_dir :value: None .. py:attribute:: best_corrcoef :value: 0 .. py:attribute:: push_to_hub :value: False .. py:attribute:: hub_model_id :value: None .. py:attribute:: hub_private_repo :value: True .. py:method:: on_epoch_end(args, state, control, **kwargs) .. py:data:: __version__ :value: '0.6.1'