Package: angle-emb
EN |简体中文AnglE📐: Angle-optimized Text EmbeddingsIt is Angle 📐, not Angel 👼.🔥A New SOTAfor Semantic Textual Similarity!🔥Our universal sentence embeddingWhereIsAI/UAE-Large-V1achieves SOTA on theMTEB Leaderboardwith an average score of 64.64!📊 Results on MTEB Leaderboard [click to expand]📊 Results on STS benchmark [click to expand]🤗 Pretrained Models🤗 HFLoRA WeightDependent BackboneLLMLanguagePromptPooling StrategyExamplesWhereIsAI/UAE-Large-V1NNNENPrompts.Cfor retrieval purposes,Nonefor othersclsSeanLee97/angle-llama-13b-nliYNousResearch/Llama-2-13b-hfYENPrompts.Alast token/SeanLee97/angle-llama-7b-nli-v2YNousResearch/Llama-2-7b-hfYENPrompts.Alast token/SeanLee97/angle-llama-7b-nli-20231027YNousResearch/Llama-2-7b-hfYENPrompts.Alast token/SeanLee97/angle-bert-base-uncased-nli-en-v1NNNENNcls_avg/SeanLee97/angle-roberta-wwm-base-zhnli-v1NNNZH-CNNcls/SeanLee97/angle-llama-7b-zhnli-v1YNousResearch/Llama-2-7b-hfYZH-CNPrompts.Blast token/💡 If the selected model is a LoRA weight, it must specify the corresponding dependent backbone.For our STS Experiment, please refer tohttps://github.com/SeanLee97/AnglE/tree/main/examples/NLIResultsEnglish STS ResultsModelSTS12STS13STS14STS15STS16STSBenchmarkSICKRelatednessAvg.SeanLee97/angle-llama-7b-nli-2023102778.6890.5885.4989.5686.9188.9281.1885.90SeanLee97/angle-llama-7b-nli-v279.0090.5685.7989.4387.0088.9780.9485.96SeanLee97/angle-llama-13b-nli79.3390.6586.8990.4587.3289.6981.3286.52SeanLee97/angle-bert-base-uncased-nli-en-v175.0985.5680.6686.4482.4785.1681.2382.37Chinese STS ResultsModelATECBQLCQMCPAWSXSTS-BSOHU-ddSOHU-dcAvg.^shibing624/text2vec-bge-large-chinese38.4161.3471.7235.1576.4471.8163.1559.72^shibing624/text2vec-base-chinese-paraphrase44.8963.5874.2440.9078.9376.7063.3063.08SeanLee97/angle-roberta-wwm-base-zhnli-v149.4972.4778.3359.1377.1472.3660.5367.06SeanLee97/angle-llama-7b-zhnli-v150.4471.9578.9056.5781.1168.1152.0265.59^ denotes baselines, their results are retrieved from:https://github.com/shibing624/text2vecUsageAnglE supports two APIs, one is thetransformersAPI, the other is theAnglEAPI. If you want to use theAnglEAPI, please install AnglE first:python-mpipinstall-Uangle-embUAEFor Retrieval PurposesFor retrieval purposes, please use the promptPrompts.C.fromangle_embimportAnglE,Promptsangle=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)For non-Retrieval Purposesfromangle_embimportAnglEangle=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)Difference between retrieval and non-retrieval sentence embeddings. [click to expand]In UAE, we use different approaches for retrieval and non-retrieval tasks, each serving a different purpose.Retrieval tasks aim to find relevant documents, and as a result, the related documents may not have strict semantic similarities to each other.For instance, when querying "How about ChatGPT?", the related documents are those that contain information related to "ChatGPT," such as "ChatGPT is amazing..." or "ChatGPT is bad....".Conversely,non-retrieval tasks, such as semantic textual similarity, require sentences that are semantically similar.For example, a sentence semantically similar to "How about ChatGPT?" could be "What is your opinion about ChatGPT?".To distinguish between these two types of tasks, we use different prompts.For retrieval tasks, we use the prompt "Represent this sentence for searching relevant passages: {text}" (Prompts.C in angle_emb).For non-retrieval tasks, we set the prompt to empty, i.e., just input your text without specifying a prompt.So, if your scenario is retrieval-related, it is highly recommended to set the prompt with angle.set_prompt(prompt=Prompts.C). If not, leave the prompt empty or use angle.set_prompt(prompt=None).Angle-LLaMAAnglEfromangle_embimportAnglE,Promptsangle=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)transformersfromangle_embimportAnglEfromtransformersimportAutoModelForCausalLM,AutoTokenizerfrompeftimportPeftModel,PeftConfigpeft_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()defdecorate_text(text:str):returnPrompts.A.format(text=text)inputs='hello world!'tok=tokenizer([decorate_text(inputs)],return_tensors='pt')fork,vintok.items():tok[k]=v.cuda()vec=model(output_hidden_states=True,**tok).hidden_states[-1][:,-1].float().detach().cpu().numpy()print(vec)Angle-BERTAnglEfromangle_embimportAnglEangle=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)transformersimporttorchfromtransformersimportAutoModel,AutoTokenizermodel_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')fork,vintok.items():tok[k]=v.cuda()hidden_state=model(**tok).last_hidden_statevec=(hidden_state[:,0]+torch.mean(hidden_state,dim=1))/2.0print(vec)Custom Train1. Data PrepationWe support two dataset formats:DatasetFormats.A: it is a pair format with three columns:text1,text2, andlabel(0/1).DatasetFormats.B: it is a triple format with three columns:text,positive, andnegative.positiveandnegativestore the positive and negative samples oftext.DatasetFormats.C: it is a pair format with two columns:text,positive.positivestore the positive sample oftext.You need to prepare your data into huggingfacedatasets.Datasetin one of the formats in terms of your supervised data.2. TrainUseangle-trainerto train your AnglE model in cli mode. Usage:CUDA_VISIBLE_DEVICES=0 angle-trainer --help3. Examplefromdatasetsimportload_datasetfromangle_embimportAnglE,AngleDataTokenizer# 1. load pretrained modelangle=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(lambdaobj:{"text1":str(obj["sentence1"]),"text2":str(obj['sentence2']),"label":obj['score']})ds=ds.select_columns(["text1","text2","label"])# 3. transform datatrain_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. fitangle.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. evaluatecorrcoef,accuracy=angle.evaluate(test_ds,device=angle.device)print('corrcoef:',corrcoef)CitationYou 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📅Description2024 Feb 7support training with only positive pairs (DatasetFormats.C)2024 Jan 11refactor to supportangle-trainerand BeLLM2023 Dec 4Release a universal English sentence embedding model:WhereIsAI/UAE-Large-V12023 Nov 2Release an English pretrained model:SeanLee97/angle-llama-13b-nli2023 Oct 28Release two chinese pretrained models:SeanLee97/angle-roberta-wwm-base-zhnli-v1andSeanLee97/angle-llama-7b-zhnli-v1; Add chinese README.md