angle_emb.angle

Classes

Prompts	Predefined prompts. Follow the model usage to choose the corresponding prompt.
AngleDataCollator	AngleDataCollator. It will be implicitly used in AnglE.fit().
Pooler	Using Pooler to obtain sentence embeddings.
AngleTrainer	Custom Huggingface Trainer for AnglE.
AngleESETrainer	Custom Huggingface Trainer for AnglE Espresso.
AngleLoss	Configure AngleLoss.
AnglE	AnglE. Everything is here👋
EvaluateCallback	Custom TrainerCallback for Angle.

Functions

detect_dataset_format(→ str)

Detect dataset format from raw data

Module Contents

angle_emb.angle.detect_dataset_format(ds: datasets.Dataset) → str

Detect dataset format from raw data

class angle_emb.angle.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)

A = 'Summarize sentence "{text}" in one word:"'

B = 'You can only output one word. Summarize "{text}":"'

C = 'Represent this sentence for searching relevant passages: {text}'

classmethod list_prompts()

class angle_emb.angle.AngleDataCollator

AngleDataCollator. It will be implicitly used in AnglE.fit(). It handles raw data, tokenizes it, and prepares batches.

Parameters:

• tokenizer – PreTrainedTokenizerBase

• padding – Union[bool, str, PaddingStrategy], padding strategy

• max_length – Optional[int], max length

• return_tensors – str

• filter_duplicate – bool. Whether filter duplicate data

• text_prompt – Optional[str], prompt for text1 and text2 (format A only). Default None

• query_prompt – Optional[str], prompt for query field. Default None

• doc_prompt – Optional[str], prompt for positive/negative fields. Default None

• dataset_format – Optional[str]. Specify dataset_format: ‘A’, ‘B’, or ‘C’. Default None.

tokenizer: transformers.tokenization_utils_base.PreTrainedTokenizerBase

padding: bool | str | transformers.utils.PaddingStrategy = 'longest'

max_length: int | None = None

return_tensors: str = 'pt'

filter_duplicate: bool = True

text_prompt: str | None = None

query_prompt: str | None = None

doc_prompt: str | None = None

dataset_format: str | None = None

static sample_from_list(text: str | List[str]) → str

Sample one string from list or return string as is

__call__(features: List[Dict], return_tensors: str = 'pt') → Dict[str, torch.Tensor]

Collate function that handles raw data.

Parameters:

• features – List[Dict]. Raw data samples

• return_tensors – str. Default “pt”

Returns:

Dict[str, torch.Tensor]. Collated data

class angle_emb.angle.Pooler(model: transformers.PreTrainedModel, pooling_strategy: int | str | None = None, padding_side: str | None = None)

Using Pooler to obtain sentence embeddings.

Parameters:

• model – PreTrainedModel

• pooling_strategy – Optional[str]. Currently support [cls, cls_avg, cls_max, last, avg, mean, max, all, int]. Default None.

• padding_side – Optional[str]. left or right. Default None.

• is_llm – bool. Default False

model

pooling_strategy = None

padding_side = None

__call__(inputs: Dict, layer_index: int = -1, embedding_start: int | None = None, embedding_size: int | None = None, return_all_layer_outputs: bool = False, pooling_strategy: int | str | None = None, return_mlm_logits: bool = False) → torch.Tensor

Get sentence embeddings.

Parameters:

• inputs – Dict. Model inputs.

• layer_index – Optional[int]. Get embeddings from specific layer.

• embedding_start – Optional[int]. Start index of embeddings.

• embedding_size – int. Set embedding size for sentence embeddings.

• return_all_layer_outputs – bool. Return all layer outputs or not. Default False.

• pooling_strategy – Optional[str]. Currently support [cls, cls_avg, cls_max, last, avg, mean, max, all, int]. Default None.

• return_mlm_logits – bool. Return logits or not. Default False.

class angle_emb.angle.AngleTrainer(pooler: Pooler, loss_kwargs: Dict | None = None, dataset_format: str = 'A', teacher_name_or_path: str | None = None, teacher_pooling_strategy: str = 'cls', pad_token_id: int = 0, model_kwargs: Dict | None = None, **kwargs)

Bases: transformers.Trainer

Custom Huggingface Trainer for AnglE.

Parameters:

• pooler – Pooler. Required

• loss_kwargs – Optional[Dict]. Default None.

• dataset_format – str. Default ‘A’

• teacher_name_or_path – Optional[str]. For distribution alignment.

• **kwargs –

  other parameters of Trainer.

pooler

pad_token_id = 0

model_kwargs = None

loss_fct

teacher_name_or_path = None

teacher_pooling_strategy = 'cls'

compute_distillation_loss(inputs: torch.Tensor, targets: torch.Tensor, mse_weight: float = 1.0, kl_temperature: float = 1.0) → torch.Tensor

Compute distillation loss.

Parameters:

• inputs – torch.Tensor. Input tensor.

• targets – torch.Tensor. Target tensor.

• mse_weight – float. MSE weight. Default 1.0.

• kl_temperature – float. KL temperature. Default 1.0.

Returns:

torch.Tensor. Distillation loss.

compute_mlm_loss(logits, mask_target_labels)

compute_loss(model, inputs, return_outputs: bool = False, **kwargs)

Compute loss for AnglE.

Parameters:

• model – Huggingface model.

• inputs – Dict. Model inputs.

• return_outputs – bool. Return outputs or not. Default False.

Returns:

torch.Tensor. Loss.

prediction_step(model, inputs, *args, **kwargs)

class angle_emb.angle.AngleESETrainer(pooler: Pooler, loss_kwargs: Dict | None = None, dataset_format: str = 'A', teacher_name_or_path: str | None = None, ese_kl_temperature: float = 1.0, ese_compression_size: int = 128, apply_ese_pca: bool = True, **kwargs)

Bases: AngleTrainer

Custom Huggingface Trainer for AnglE Espresso.

Parameters:

• pooler – Pooler. Required

• loss_kwargs – Optional[Dict]. Default None.

• dataset_format – str. Default ‘A’

• teacher_name_or_path – Optional[str]. For distribution alignment.

• **kwargs –

  other parameters of Trainer.

ese_kl_temperature = 1.0

ese_compression_size = 128

apply_ese_pca = True

n_layers

pca_compress(m: torch.Tensor, k: int) → torch.Tensor

Get topk feature via PCA.

Parameters:

• m – torch.Tensor. Input tensor.

• k – int. Top-k feature size.

Returns:

torch.Tensor. Top-k feature.

compute_student_loss(inputs: Dict, all_layer_outputs: torch.Tensor, labels: torch.Tensor, pooling_strategy: str, padding_side: str) → torch.Tensor

compute_loss(model, inputs, return_outputs=False, **kwargs)

Compute loss for Espresso.

Parameters:

• model – Huggingface model.

• inputs – Dict. Model inputs.

• return_outputs – bool. Return outputs or not. Default False.

Returns:

torch.Tensor. Loss.

class angle_emb.angle.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: str | None = None, **kwargs)

Configure AngleLoss.

Parameters:

• cosine_w – float. weight for cosine_loss. Default 1.0

• ibn_w – float. weight for in batch negative loss. Default 1.0

• cln_w – float. weight for contrastive learning with hard negative. Default 1.0

• angle_w – float. weight for angle loss. Default 1.0

• cosine_tau – float. tau for cosine loss. Default 20.0

• ibn_tau – float. tau for in batch negative loss. Default 20.0

• angle_tau – float. tau for angle loss. Default 20.0

• angle_pooling_strategy – str. pooling strategy for angle loss. Default’sum’.

• dataset_format – Optional[str]. Default None.

cosine_w = 0.0

ibn_w = 1.0

cln_w = 1.0

angle_w = 0.02

cosine_tau = 20.0

ibn_tau = 20.0

angle_tau = 20.0

angle_pooling_strategy = 'sum'

dataset_format = None

__call__(labels: torch.Tensor, outputs: torch.Tensor) → torch.Tensor

Compute loss for AnglE.

Parameters:

• labels – torch.Tensor. Labels.

• outputs – torch.Tensor. Outputs.

Returns:

torch.Tensor. Loss.

class angle_emb.angle.AnglE(model_name_or_path: str, tokenizer_name_or_path: str | None = None, max_length: int = 512, model_kwargs: Dict | None = None, lora_config_kwargs: Dict | None = None, pooling_strategy: str | None = None, apply_lora: bool | None = None, train_mode: bool = True, load_kbit: int | None = None, is_llm: bool | None = None, pretrained_model_path: str | None = None, pretrained_lora_path: str | None = None, torch_dtype: torch.dtype | None = None, device: str | None = None, kbit_kwargs: Dict | None = None, tokenizer_padding_side: str | None = None, apply_billm: bool = False, billm_model_class: str | None = None, load_mlm_model: bool = False, **kwargs: Any)

Bases: angle_emb.base.AngleBase

AnglE. Everything is here👋

Parameters:

• model_name_or_path – str, model name or path.

• tokenizer_name_or_path – Optional[str]. Default None. When it set to None, it will use the same as model_name_or_path.

• max_length – int. Default 512

• model_kwargs – Optional[Dict]. kwargs for model.

• lora_config_kwargs – Optional[Dict]. kwargs for peft lora_config. details refer to: https://huggingface.co/docs/peft/tutorial/peft_model_config

• pooling_strategy – Optional[str]. Pooling strategy. Currently support [cls, cls_avg, cls_max, last, avg, mean, max, all, int]

• apply_lora – Optional[bool]. Whether apply lora. Default None.

• train_mode – bool. Whether load for training. Default True.

• load_kbit – Optional[int]. Specify kbit training from [4, 8, 16]. Default None.

• is_llm – Optional[bool]. Whether the model is llm. Must be set manually. Default None.

• pretrained_model_path – Optional[str]. Default None.

• pretrained_lora_path – Optional[str]. Default None.

• torch_dtype – Optional[torch.dtype]. Specify torch_dtype. Default None.

• device – Optional[str]. Specify device. Default None.

• 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

• tokenizer_padding_side – Optional[str]. Specify tokenizer padding side from [left, right]. Default None.

• apply_billm – bool. Whether apply billm. Default False.

• billm_model_class – Optional[str]. Specify billm model class. Default None.

• load_mlm_model – bool. Whether load mlm model. Default False. If set True, it will load model with AutoModelForMaskedLM.

• **kwargs –

  Any.

cfg_file_name = 'angle_config.json'

special_columns = ['labels']

max_length = 512

train_mode = True

pooling_strategy = None

load_kbit = None

is_llm = None

load_mlm_model = False

apply_lora = None

model_kwargs

tokenizer

pooler

__cfg

cuda()

to(device: Any)

static kbit_post_handle(model: torch.nn.Module) → torch.nn.Module

static find_pth_path(dirpath: str, config: Dict) → str

static from_pretrained(model_name_or_path: str, pretrained_model_path: str | None = None, pretrained_lora_path: str | None = None, is_llm: bool | None = None, pooling_strategy: str = 'cls', train_mode: bool = False, model_kwargs: Dict | None = None, **kwargs)

Load AnglE from pretrained model.

Parameters:

• model_name_or_path – str, model name or path. Required.

• pretrained_model_path – Optional[str].

• pretrained_lora_path – Optional[str].

• is_llm – Optional[bool].

• pooling_strategy – str. Pooling Strategy. Default cls.

• train_mode – bool. Default False.

• kwargs – Other kwargs for AnglE.

Returns:

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)

static load_config(fpath: str) → Dict

save_config(fpath: str)

fit(train_ds: datasets.Dataset, valid_ds: datasets.Dataset | None = None, valid_ds_for_callback: datasets.Dataset | None = None, batch_size: int = 32, output_dir: str | None = 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: bool | None = None, bf16: bool | None = None, argument_kwargs: Dict | None = None, trainer_kwargs: Dict | None = None, loss_kwargs: Dict | None = None, apply_ese: bool = False, filter_duplicate: bool = True, push_to_hub: bool = False, hub_model_id: str | None = None, hub_private_repo: bool = True, padding: str = 'longest', text_prompt: str | None = None, query_prompt: str | None = None, doc_prompt: str | None = None)

Fit using AnglE.

Parameters:

• train_ds – Dataset. Raw train dataset (not tokenized). Required.

• valid_ds – Optional[Dataset]. Raw valid dataset (not tokenized). Default None.

• 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.

• batch_size – int. Default 32.

• output_dir – Optional[str]. save dir. Default None.

• epochs – int. Default 1.

• learning_rate – float. Default 1e-5.

• warmup_steps – int. Default 1000.

• logging_steps – int. Default 10.

• eval_steps – int. Default 1000.

• eval_strategy – str. Default ‘steps’.

• save_steps – int. Default 100.

• save_strategy – str. Default steps.

• save_total_limit – int. Default 10.

• gradient_accumulation_steps – int. Default 1.

• fp16 – Optional[bool]. Default None.

• bf16 – Optional[bool]. Default None.

• argument_kwargs – Optional[Dict]. kwargs for TrainingArguments. refer to: https://huggingface.co/docs/transformers/v4.37.0/en/main_classes/trainer#transformers.TrainingArguments

• trainer_kwargs – Optional[Dict]. kwargs for AngleTrainer.

• loss_kwargs – Optional[Dict]. kwargs for AngleLoss.

• apply_ese – bool, whether apply ESE training.

• filter_duplicate – bool, whether filter duplicate samples.

• push_to_hub – bool, whether push to hub.

• hub_model_id – Optional[str], hub model id.

• hub_private_repo – bool, whether push to private repo.

• padding – str, padding strategy of tokenizer. Default ‘longest’.

• text_prompt – Optional[str], prompt for text1 and text2 (format A only). Default None.

• query_prompt – Optional[str], prompt template for query. Default None.

• doc_prompt – Optional[str], prompt template for documents (positive/negative). Default None.

evaluate(ds: datasets.Dataset, batch_size: int = 32, metric: str = 'spearman_cosine', prompt: str | None = None) → float

evaluate

Parameters:

• data – Dataset, format A is required

• batch_size – int. Default 32.

• metric – str. Default ‘spearman_cosine’.

Returns:

float.

truncate_layer(layer_index: int)

truncate layer

Parameters:

layer_index – int. layers after layer_index will be truncated.

Returns:

self

encode(inputs: List[str] | Tuple[str] | str, max_length: int | None = None, to_numpy: bool = True, embedding_start: int = 0, embedding_size: int | None = None, device: Any | None = None, prompt: str | None = None, normalize_embedding: bool = False, padding: str = 'longest')

encode texts.

Parameters:

• inputs – Union[List[str], Tuple[str], List[Dict], str]. Input texts. Required.

• max_length – Optional[int]. Default None.

• to_numpy – bool. Default True.

• embedding_start – int. Specify the start position of the embedding (for Espresso).

• embedding_size – Optional[int]. Specify embedding size (for Espresso). The embeddings from embedding_start to embedding_start+embedding_size will be returned.

• device – Optional[Any]. Default None.

• prompt – Optional[str]. Default None.

• normalize_embedding – bool. Default False.

• padding – str. Padding strategy of tokenizer. Default ‘longest’.

push_to_hub(hub_model_id: str, private: bool = True, exist_ok: bool = False, **kwargs)

push model to hub

Parameters:

• hub_model_id – str, hub model id.

• private – bool, whether push to private repo. Default True.

• exist_ok – bool, whether allow overwrite. Default False.

• kwargs – other kwargs for push_to_hub method.

save_pretrained(output_dir: str, exist_ok: bool = True)

save model and tokenizer

Parameters:

• output_dir – str, output dir.

• exist_ok – bool, whether allow overwrite. Default True.

class angle_emb.angle.EvaluateCallback(model: AnglE, valid_ds: datasets.Dataset, evaluate_fn: Callable, save_dir: str | None = None, push_to_hub: bool = False, hub_model_id: str | None = None, hub_private_repo: bool = True)

Bases: transformers.TrainerCallback

Custom TrainerCallback for Angle. This callback will compute corrcoef for each epoch.

Parameters:

• model – PreTrainedModel.

• valid_ds – Dataset.

• evaluate_fn – Callable. It will receive valid_ds as input like evaluate_fn(valid_ds).

• save_dir – Optional[str]. specify dir to save model with best results.

model

valid_ds

evaluate_fn

save_dir = None

best_corrcoef = 0

push_to_hub = False

hub_model_id = None

hub_private_repo = True

on_epoch_end(args, state, control, **kwargs)