EN | 简体中文

AnglE📐: Angle-optimized Text Embeddings

It is Angle 📐, not Angel 👼.

🔥 A New SOTA for Semantic Textual Similarity!

🔥 Our universal sentence embedding WhereIsAI/UAE-Large-V1 achieves SOTA on the MTEB Leaderboard with an average score of 64.64!

🤗 Pretrained Models

🤗 HF	LoRA Weight	Dependent Backbone	LLM	Language	Prompt	Pooling Strategy	Examples
WhereIsAI/UAE-Large-V1	N	N	N	EN	Prompts.C for retrieval purposes, None for others	cls
SeanLee97/angle-llama-13b-nli	Y	NousResearch/Llama-2-13b-hf	Y	EN	Prompts.A	last token	/
SeanLee97/angle-llama-7b-nli-v2	Y	NousResearch/Llama-2-7b-hf	Y	EN	Prompts.A	last token	/
SeanLee97/angle-llama-7b-nli-20231027	Y	NousResearch/Llama-2-7b-hf	Y	EN	Prompts.A	last token	/
SeanLee97/angle-bert-base-uncased-nli-en-v1	N	N	N	EN	N	cls_avg	/
SeanLee97/angle-roberta-wwm-base-zhnli-v1	N	N	N	ZH-CN	N	cls	/
SeanLee97/angle-llama-7b-zhnli-v1	Y	NousResearch/Llama-2-7b-hf	Y	ZH-CN	Prompts.B	last token	/

💡 If the selected model is a LoRA weight, it must specify the corresponding dependent backbone.

For our STS Experiment, please refer to https://github.com/SeanLee97/AnglE/tree/main/examples/NLI

Results

English STS Results

Chinese STS Results

^ denotes baselines, their results are retrieved from: https://github.com/shibing624/text2vec

Usage

AnglE supports two APIs, one is the transformers API, the other is the AnglE API. If you want to use the AnglE API, please install AnglE first:

python -m pip install -U angle-emb

UAE

1. For Retrieval Purposes

For retrieval purposes, please use the prompt Prompts.C for the query (⚠️：no need to apply prompt for documents).

from angle_emb import AnglE, Prompts

angle = AnglE.from_pretrained('WhereIsAI/UAE-Large-V1', pooling_strategy='cls').cuda()
angle.set_prompt(prompt=Prompts.C)
vec = angle.encode({'text': 'hello world'}, to_numpy=True)
print(vec)
vecs = angle.encode([{'text': 'hello world1'}, {'text': 'hello world2'}], to_numpy=True)
print(vecs)

2. For non-Retrieval Purposes

from angle_emb import AnglE

angle = AnglE.from_pretrained('WhereIsAI/UAE-Large-V1', pooling_strategy='cls').cuda()
vec = angle.encode('hello world', to_numpy=True)
print(vec)
vecs = angle.encode(['hello world1', 'hello world2'], to_numpy=True)
print(vecs)

Angle-LLaMA

1. AnglE

from angle_emb import AnglE, Prompts

angle = AnglE.from_pretrained('NousResearch/Llama-2-7b-hf', pretrained_lora_path='SeanLee97/angle-llama-7b-nli-v2')

print('All predefined prompts:', Prompts.list_prompts())
angle.set_prompt(prompt=Prompts.A)
print('prompt:', angle.prompt)
vec = angle.encode({'text': 'hello world'}, to_numpy=True)
print(vec)
vecs = angle.encode([{'text': 'hello world1'}, {'text': 'hello world2'}], to_numpy=True)
print(vecs)

2. transformers

from angle_emb import AnglE, Prompts
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel, PeftConfig

peft_model_id = 'SeanLee97/angle-llama-7b-nli-v2'
config = PeftConfig.from_pretrained(peft_model_id)
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path).bfloat16().cuda()
model = PeftModel.from_pretrained(model, peft_model_id).cuda()

def decorate_text(text: str):
    return Prompts.A.format(text=text)

inputs = 'hello world!'
tok = tokenizer([decorate_text(inputs)], return_tensors='pt')
for k, v in tok.items():
    tok[k] = v.cuda()
vec = model(output_hidden_states=True, **tok).hidden_states[-1][:, -1].float().detach().cpu().numpy()
print(vec)

Angle-BERT

1. AnglE

from angle_emb import AnglE

angle = AnglE.from_pretrained('SeanLee97/angle-bert-base-uncased-nli-en-v1', pooling_strategy='cls_avg').cuda()
vec = angle.encode('hello world', to_numpy=True)
print(vec)
vecs = angle.encode(['hello world1', 'hello world2'], to_numpy=True)
print(vecs)

2. transformers

import torch
from transformers import AutoModel, AutoTokenizer

model_id = 'SeanLee97/angle-bert-base-uncased-nli-en-v1'
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModel.from_pretrained(model_id).cuda()

inputs = 'hello world!'
tok = tokenizer([inputs], return_tensors='pt')
for k, v in tok.items():
    tok[k] = v.cuda()
hidden_state = model(**tok).last_hidden_state
vec = (hidden_state[:, 0] + torch.mean(hidden_state, dim=1)) / 2.0
print(vec)

Custom Train

1. Data Prepation

We support two dataset formats:

1. DatasetFormats.A: it is a pair format with three columns: text1, text2, and label (0/1).

2. DatasetFormats.B: it is a triple format with three columns: text, positive, and negative. positive and negative store the positive and negative samples of text.

3. DatasetFormats.C: it is a pair format with two columns: text, positive. positive store the positive sample of text.

You need to prepare your data into huggingface datasets.Dataset in one of the formats in terms of your supervised data.

2. Train

Use angle-trainer to train your AnglE model in cli mode. Usage: CUDA_VISIBLE_DEVICES=0 angle-trainer --help

3. Example

from datasets import load_dataset
from angle_emb import AnglE, AngleDataTokenizer


# 1. load pretrained model
angle = AnglE.from_pretrained('SeanLee97/angle-bert-base-uncased-nli-en-v1', max_length=128, pooling_strategy='cls').cuda()

# 2. load dataset
# `text1`, `text2`, and `label` are three required columns.
ds = load_dataset('mteb/stsbenchmark-sts')
ds = ds.map(lambda obj: {"text1": str(obj["sentence1"]), "text2": str(obj['sentence2']), "label": obj['score']})
ds = ds.select_columns(["text1", "text2", "label"])

# 3. transform data
train_ds = ds['train'].shuffle().map(AngleDataTokenizer(angle.tokenizer, angle.max_length), num_proc=8)
valid_ds = ds['validation'].map(AngleDataTokenizer(angle.tokenizer, angle.max_length), num_proc=8)
test_ds = ds['test'].map(AngleDataTokenizer(angle.tokenizer, angle.max_length), num_proc=8)

# 4. fit
angle.fit(
    train_ds=train_ds,
    valid_ds=valid_ds,
    output_dir='ckpts/sts-b',
    batch_size=32,
    epochs=5,
    learning_rate=2e-5,
    save_steps=100,
    eval_steps=1000,
    warmup_steps=0,
    gradient_accumulation_steps=1,
    loss_kwargs={
        'w1': 1.0,
        'w2': 1.0,
        'w3': 1.0,
        'cosine_tau': 20,
        'ibn_tau': 20,
        'angle_tau': 1.0
    },
    fp16=True,
    logging_steps=100
)

# 5. evaluate
corrcoef, accuracy = angle.evaluate(test_ds, device=angle.device)
print('corrcoef:', corrcoef)

4. Fine-tuning Tips 💡

1. if your dataset format is DatasetFormats.A, it is recommended to slightly increase the weight for w1 or slightly decrease the weight for w2.

2. if your dataset format is DatasetFormats.B, it is recommended to set w1 to 0, and increase the weight for w2 such as 10 and 20. The angle_tau can be set to 20.0.

3. if your dataset format is DatasetFormats.C, only w2 and ibn_tau are effective. You don't need to tune other parameters.

4. To alleviate information forgetting in fine-tuning, it is better to specify the fixed_teacher_name_or_path. If the fixed_teacher_name_or_path equals model_name_or_path, it will conduct self-distillation. It is worth to note that fixed_teacher_name_or_path has to have the same tokenizer as model_name_or_path. Or it will lead to unexpected results.

Citation

You are welcome to use our code and pre-trained models. If you use our code and pre-trained models, please support us by citing our work as follows:

@article{li2023angle,
  title={AnglE-optimized Text Embeddings},
  author={Li, Xianming and Li, Jing},
  journal={arXiv preprint arXiv:2309.12871},
  year={2023}
}

ChangeLogs