angle_emb.angle =============== .. py:module:: angle_emb.angle Classes ------- .. autoapisummary:: angle_emb.angle.Prompts angle_emb.angle.AngleDataCollator angle_emb.angle.Pooler angle_emb.angle.AngleTrainer angle_emb.angle.AngleESETrainer angle_emb.angle.AngleLoss angle_emb.angle.AnglE angle_emb.angle.EvaluateCallback Functions --------- .. autoapisummary:: angle_emb.angle.detect_dataset_format Module Contents --------------- .. 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)