语言:

• 中文
• 英文 标签:
• 句子转换器
• 句子相似度
• 特征提取
• 转换器 管道标签: 句子相似度 库名称: 句子转换器 许可证: apache-2.0

FlagEmbedding

更多详情请参考我们的Github: FlagEmbedding.

BGE-Code-v1 是一个基于LLM的代码嵌入模型，支持代码检索、文本检索和多语言检索。它主要展示了以下能力：

• 卓越的代码检索性能：该模型展示了出色的代码检索能力，支持中英文自然语言查询以及20种编程语言。
• 强大的文本检索能力：该模型保持了与同类规模文本嵌入模型相当的强大文本检索能力。
• 广泛的多语言支持：BGE-Code-v1提供了全面的多语言检索能力，在英语、中文、日语、法语等语言上表现优异。

使用方法

使用FlagEmbedding

git clone https://github.com/FlagOpen/FlagEmbedding.git
cd FlagEmbedding
pip install -e .

from FlagEmbedding import FlagLLMModel
queries = [
    "从'Staff'表中删除ID为4的记录。", 
    '删除"Livestock"表中年龄大于5的所有记录'
]
documents = [
    "DELETE FROM Staff WHERE StaffID = 4;",
    "DELETE FROM Livestock WHERE age > 5;"
]
model = FlagLLMModel('BAAI/bge-code-v1', 
                     query_instruction_format="<instruct>{}\n<query>{}",
                     query_instruction_for_retrieval="给定一个文本问题，检索适合回答该问题的SQL查询。",
                     trust_remote_code=True,
                     use_fp16=True) # 设置use_fp16为True可以加速计算，但性能略有下降
embeddings_1 = model.encode_queries(queries)
embeddings_2 = model.encode_corpus(documents)
similarity = embeddings_1 @ embeddings_2.T
print(similarity)

默认情况下，FlagLLMModel在编码时会使用所有可用的GPU。请设置os.environ["CUDA_VISIBLE_DEVICES"]来选择特定的GPU。你也可以设置os.environ["CUDA_VISIBLE_DEVICES"]=""来禁用所有GPU。

使用Sentence Transformers

from sentence_transformers import SentenceTransformer
import torch

# 加载模型，可以选择使用float16精度以加快推理速度
model = SentenceTransformer(
    "BAAI/bge-code-v1",
    trust_remote_code=True,
    model_kwargs={"torch_dtype": torch.float16},
)

# 根据指令准备提示
instruction = '给定一个文本问题，检索适合回答该问题的SQL查询。'
prompt = f'<instruct>{instruction}\n<query>'
# 准备查询和文档
queries = [
    "从'Staff'表中删除ID为4的记录。", 
    '删除"Livestock"表中年龄大于5的所有记录'
]
documents = [
    "DELETE FROM Staff WHERE StaffID = 4;",
    "DELETE FROM Livestock WHERE age > 5;"
]

# 计算查询和文档的嵌入
query_embeddings = model.encode(queries, prompt=prompt)
document_embeddings = model.encode(documents)

# 计算查询和文档嵌入之间的余弦相似度
similarities = model.similarity(query_embeddings, document_embeddings)
print(similarities)

使用HuggingFace Transformers

import torch
import torch.nn.functional as F

from torch import Tensor
from transformers import AutoTokenizer, AutoModel


def last_token_pool(last_hidden_states: Tensor,
                 attention_mask: Tensor) -> Tensor:
    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
    if left_padding:
        return last_hidden_states[:, -1]
    else:
        sequence_lengths = attention_mask.sum(dim=1) - 1
        batch_size = last_hidden_states.shape[0]
        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]


def get_detailed_instruct(task_description: str, query: str) -> str:
    return f'<instruct>{task_description}\n<query>{query}'


instruction = '给定一个文本问题，检索适合回答该问题的SQL查询。'
queries = [
    "从'Staff'表中删除ID为4的记录。", 
    '删除"Livestock"表中年龄大于5的所有记录'
]
documents = [
    "DELETE FROM Staff WHERE StaffID = 4;",
    "DELETE FROM Livestock WHERE age > 5;"
]
input_texts = queries + documents

tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-code-v1', trust_remote_code=True)
model = AutoModel.from_pretrained('BAAI/bge-code-v1', trust_remote_code=True)
model.eval()

max_length = 4096
# 对输入文本进行分词
batch_dict = tokenizer(input_texts, max_length=max_length, padding=True, truncation=True, return_tensors='pt', pad_to_multiple_of=8)

with torch.no_grad():
    outputs = model(**batch_dict)
    embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
    
# 归一化嵌入
embeddings = F.normalize(embeddings, p=2, dim=1)
scores = (embeddings[:2] @ embeddings[2:].T) * 100
print(scores.tolist())

评估

BGE-Code-v1 在CoIR和CodeRAG基准测试中均达到了最先进的性能。

• CoIR

• CodedRAG

评估指令

{
    "Apps": "给定一个代码竞赛问题描述，检索有助于解决问题的相关代码。",
    "CosQA": "给定一个网页搜索查询，检索有助于回答查询的相关代码。",
    "Text2SQL": "给定一个文本问题，检索适合回答该问题的SQL查询。",
    "CSN": "给定一段代码，检索总结该代码的文档字符串。",
    "CSN-CCR": "给定一段代码片段，检索作为该代码后续部分的代码片段。",
    "CodeTrans-DL": "给定一段代码，检索与输入代码语义上等效的代码。",
    "CodeTrans-Contest": "给定一段Python代码，检索与输入代码语义上等效的C++代码。",
    "StackOverFlow-QA": "给定一个包含文本和代码片段混合的问题，检索包含文本和代码片段混合的相关答案，并能帮助回答问题。",
    "CodeFeedBack-ST": "给定一个包含文本和代码片段混合的问题，检索包含文本和代码片段混合的相关答案，并能帮助回答问题。",
    "CodeFeedBack-MT": "给定一个包含文本和代码片段混合的多轮对话历史，检索包含文本和代码片段混合的相关答案，并能帮助回答问题。",
    "HummanEval": "给定一个包含文本和代码片段混合的问题，检索包含文本和代码片段混合的相关答案，并能帮助回答问题。",
    "MBPP": "给定一个代码功能的文本解释，检索对应的代码实现。",
    "DS-1000": "给定一个包含文本和代码片段混合的问题，检索包含文本和代码片段混合的相关答案，并能帮助回答问题。",
    "ODEX": "给定一个问题，检索包含文本和代码片段混合的相关答案，并能帮助回答问题。",
    "RepoEval": "给定一段代码片段，检索作为该代码后续部分的代码片段。",
    "SWE-bench-Lite": "给定一个包含错误和错误或错误的自然语言描述的代码片段，检索展示类似错误或错误的解决方案或修复的代码片段（期望的文档）。"
}

引用

如果您觉得这个仓库有用，请考虑给予星标 :star: 和引用

@article{bge-llm,
  title={Making text embedders few-shot learners},
  author={Li, Chaofan and Qin, MingHao and Xiao, Shitao and Chen, Jianlyu and Luo, Kun and Shao, Yingxia and Lian, Defu and Liu, Zheng},
  journal={arXiv preprint arXiv:2409.15700},
  year={2024}
}

@misc{bge-m3,
      title={BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation}, 
      author={Jianlv Chen and Shitao Xiao and Peitian Zhang and Kun Luo and Defu Lian and Zheng Liu},
      year={2024},
      eprint={2402.03216},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}

@misc{bge_embedding,
      title={C-Pack: Packaged Resources To Advance General Chinese Embedding}, 
      author={Shitao Xiao and Zheng Liu and Peitian Zhang and Niklas Muennighoff},
      year={2023},
      eprint={2309.07597},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}