angle_emb

Submodules

Attributes

`logger`
`__version__`
`__version__`

Classes

`AngleBase`
`CorrelationEvaluator`
`AngleBase`
`CorrelationEvaluator`
`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

`categorical_crossentropy_loss`(→ torch.Tensor)	Compute categorical crossentropy
`cosine_loss`(→ torch.Tensor)	Compute cosine loss
`angle_loss`(y_true, y_pred[, tau, pooling_strategy])	Compute angle loss
`in_batch_negative_loss`(→ torch.Tensor)	Compute in-batch negative loss, i.e., contrastive loss
`contrastive_with_negative_loss`(→ torch.Tensor)	Compute contrastive with negative loss
`angle_loss`(y_true, y_pred[, tau, pooling_strategy])	Compute angle loss
`contrastive_with_negative_loss`(→ torch.Tensor)	Compute contrastive with negative loss
`cosine_loss`(→ torch.Tensor)	Compute cosine loss
`in_batch_negative_loss`(→ torch.Tensor)	Compute in-batch negative loss, i.e., contrastive loss
`find_all_linear_names`(→ List[str])	Find all linear layer names
`get_pooling`(→ torch.Tensor)	Pooling the model outputs.
`set_device`(→ str)	Set device automatically
`detect_dataset_format`(→ str)	Detect dataset format from raw data

Package Contents

angle_emb.categorical_crossentropy_loss(y_true: torch.Tensor, y_pred: torch.Tensor, from_logits: bool = True) → torch.Tensor[source]

Compute categorical crossentropy

Parameters:

y_true – torch.Tensor, ground truth
y_pred – torch.Tensor, model output
from_logits – bool, True means y_pred has not transformed by softmax, default True

Returns:

torch.Tensor, loss value

angle_emb.cosine_loss(y_true: torch.Tensor, y_pred: torch.Tensor, tau: float = 20.0) → torch.Tensor[source]

Compute cosine loss

Parameters:

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.
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.
tau – float, scale factor, default 20

Returns:

torch.Tensor, loss value

angle_emb.angle_loss(y_true: torch.Tensor, y_pred: torch.Tensor, tau: float = 1.0, pooling_strategy: str = 'sum')[source]

Compute angle loss

Parameters:

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.
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.
tau – float, scale factor, default 1.0

Returns:

torch.Tensor, loss value

angle_emb.in_batch_negative_loss(y_true: torch.Tensor, y_pred: torch.Tensor, tau: float = 20.0, negative_weights: float = 0.0) → torch.Tensor[source]

Compute in-batch negative loss, i.e., contrastive loss

Parameters:

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.
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.
tau – float, scale factor, default 20.0
negative_weights – float, negative weights, default 0.0

Returns:

torch.Tensor, loss value

angle_emb.contrastive_with_negative_loss(text: torch.Tensor, pos: torch.Tensor, neg: torch.Tensor | None = None, tau: float = 20.0) → torch.Tensor[source]

Compute contrastive with negative loss

Parameters:

text – torch.Tensor, text.
pos – torch.Tensor, positive samples of text.
neg – torch.Tensor, negative samples of text.
tau – float, scale factor, default 20.0

Returns:

torch.Tensor, loss value

class angle_emb.AngleBase[source]

abstractmethod encode()[source]

abstractmethod fit()[source]

class angle_emb.CorrelationEvaluator(text1: List[str], text2: List[str], labels: List[float], batch_size: int = 32)[source]

Bases: object

text1

text2

labels

batch_size = 32

__call__(model: angle_emb.base.AngleBase, show_progress: bool = True, **kwargs) → dict[source]

Evaluate the model on the given dataset.

Parameters:

model – AnglE, the model to evaluate.
show_progress – bool, whether to show a progress bar during evaluation.
kwargs – Additional keyword arguments to pass to the encode method of the model.

Returns:

dict, The evaluation results.

list_all_metrics() → List[str][source]

Get a list of all the metrics that can be computed by this evaluator.

Returns:: List[str], A list of all the metrics that can be computed by this evaluator.

class angle_emb.AngleBase[source]

abstractmethod encode()[source]

abstractmethod fit()[source]

class angle_emb.CorrelationEvaluator(text1: List[str], text2: List[str], labels: List[float], batch_size: int = 32)[source]

Bases: object

text1

text2

labels

batch_size = 32

__call__(model: angle_emb.base.AngleBase, show_progress: bool = True, **kwargs) → dict[source]

Evaluate the model on the given dataset.

Parameters:

model – AnglE, the model to evaluate.
show_progress – bool, whether to show a progress bar during evaluation.
kwargs – Additional keyword arguments to pass to the encode method of the model.

Returns:

dict, The evaluation results.

list_all_metrics() → List[str][source]

Get a list of all the metrics that can be computed by this evaluator.

Returns:: List[str], A list of all the metrics that can be computed by this evaluator.

angle_emb.angle_loss(y_true: torch.Tensor, y_pred: torch.Tensor, tau: float = 1.0, pooling_strategy: str = 'sum')[source]

Compute angle loss

Parameters:

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.
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.
tau – float, scale factor, default 1.0

Returns:

torch.Tensor, loss value

angle_emb.contrastive_with_negative_loss(text: torch.Tensor, pos: torch.Tensor, neg: torch.Tensor | None = None, tau: float = 20.0) → torch.Tensor[source]

Compute contrastive with negative loss

Parameters:

text – torch.Tensor, text.
pos – torch.Tensor, positive samples of text.
neg – torch.Tensor, negative samples of text.
tau – float, scale factor, default 20.0

Returns:

torch.Tensor, loss value

angle_emb.cosine_loss(y_true: torch.Tensor, y_pred: torch.Tensor, tau: float = 20.0) → torch.Tensor[source]

Compute cosine loss

Parameters:

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.
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.
tau – float, scale factor, default 20

Returns:

torch.Tensor, loss value

angle_emb.in_batch_negative_loss(y_true: torch.Tensor, y_pred: torch.Tensor, tau: float = 20.0, negative_weights: float = 0.0) → torch.Tensor[source]

Compute in-batch negative loss, i.e., contrastive loss

Parameters:

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.
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.
tau – float, scale factor, default 20.0
negative_weights – float, negative weights, default 0.0

Returns:

torch.Tensor, loss value

angle_emb.find_all_linear_names(model: transformers.PreTrainedModel, linear_type: object | None = None) → List[str][source]

Find all linear layer names

Parameters:

model – PreTrainedModel
linear_type – Optional[object] = None, linear type, such as nn.Linear and bnb.nn.Linear4bit.

Returns:

List[str], linear layer names

angle_emb.get_pooling(outputs: torch.Tensor, inputs: Dict, pooling_strategy: str, padding_side: str) → torch.Tensor[source]

Pooling the model outputs.

Parameters:

outputs – torch.Tensor. Model outputs (without pooling)
inputs – Dict. Model inputs
pooling_strategy – str. Pooling strategy [cls, cls_avg, cls_max, last, avg, mean, max, all, int]
padding_side – str. Padding strategy of tokenizers (left or right). It can be obtained by tokenizer.padding_side.

angle_emb.logger[source]

angle_emb.set_device() → str[source]

Set device automatically

Returns:: str, device name

angle_emb.__version__ = '0.6.1'[source]

angle_emb.detect_dataset_format(ds: datasets.Dataset) → str[source]: Detect dataset format from raw data

class angle_emb.Prompts[source]

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()[source]

class angle_emb.AngleDataCollator[source]

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[source]: Sample one string from list or return string as is

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

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.Pooler(model: transformers.PreTrainedModel, pooling_strategy: int | str | None = None, padding_side: str | None = None)[source]

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[source]

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.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)[source]

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[source]

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)[source]

compute_loss(model, inputs, return_outputs: bool = False, **kwargs)[source]

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)[source]

class angle_emb.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)[source]

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[source]

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[source]

compute_loss(model, inputs, return_outputs=False, **kwargs)[source]

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.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)[source]

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[source]

Compute loss for AnglE.

Parameters:

labels – torch.Tensor. Labels.
outputs – torch.Tensor. Outputs.

Returns:

torch.Tensor. Loss.

class angle_emb.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)[source]

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()[source]

to(device: Any)[source]

static kbit_post_handle(model: torch.nn.Module) → torch.nn.Module[source]

static find_pth_path(dirpath: str, config: Dict) → str[source]

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)[source]

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[source]

save_config(fpath: str)[source]

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)[source]

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[source]

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)[source]

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')[source]

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)[source]

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)[source]

save model and tokenizer

Parameters:

output_dir – str, output dir.
exist_ok – bool, whether allow overwrite. Default True.

class angle_emb.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)[source]

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)[source]

angle_emb.__version__ = '0.6.1'[source]