BAAI/bge-code-v1
- License
- apache-2.0
- Size
- not mirrored yet
No official torrent yet — requesting it tells the flock what to host next. You cannot seed a catalog-only page.
Cataloged from BAAI/bge-code-v1
All models by BAAI →No official torrent yet — requesting it tells the flock what to host next. You cannot seed a catalog-only page.
Cataloged from BAAI/bge-code-v1
All models by BAAI →For more details please refer to our Github: FlagEmbedding.
BGE-Code-v1 is an LLM-based code embedding model that supports code retrieval, text retrieval, and multilingual retrieval. It primarily demonstrates the following capabilities:
By default, FlagLLMModel will use all available GPUs when encoding. Please set os.environ["CUDA_VISIBLE_DEVICES"] to select specific GPUs. You also can set os.environ["CUDA_VISIBLE_DEVICES"]="" to make all GPUs unavailable.
BGE-Code-v1 achieves state-of-the-art performance on both the CoIR and CodeRAG benchmarks.
If you find this repository useful, please consider giving a star :star: and citation
git clone https://github.com/FlagOpen/FlagEmbedding.gitcd FlagEmbeddingpip install -e .from FlagEmbedding import FlagLLMModelqueries = [ "Delete the record with ID 4 from the 'Staff' table.", 'Delete all records in the "Livestock" table where age is greater than 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="Given a question in text, retrieve SQL queries that are appropriate responses to the question.", trust_remote_code=True, use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradationembeddings_1 = model.encode_queries(queries)embeddings_2 = model.encode_corpus(documents)similarity = embeddings_1 @ embeddings_2.Tprint(similarity)from sentence_transformers import SentenceTransformerimport torch
# Load the model, optionally in float16 precision for faster inferencemodel = SentenceTransformer( "BAAI/bge-code-v1", trust_remote_code=True, model_kwargs={"torch_dtype": torch.float16},)
# Prepare a prompt given an instructioninstruction = 'Given a question in text, retrieve SQL queries that are appropriate responses to the question.'prompt = f'<instruct>{instruction}\n<query>'# Prepare queries and documentsqueries = [ "Delete the record with ID 4 from the 'Staff' table.", 'Delete all records in the "Livestock" table where age is greater than 5']documents = [ "DELETE FROM Staff WHERE StaffID = 4;", "DELETE FROM Livestock WHERE age > 5;"]
# Compute the query and document embeddingsquery_embeddings = model.encode(queries, prompt=prompt)document_embeddings = model.encode(documents)
# Compute the cosine similarity between the query and document embeddingssimilarities = model.similarity(query_embeddings, document_embeddings)print(similarities)import torchimport torch.nn.functional as F
from torch import Tensorfrom 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 = 'Given a question in text, retrieve SQL queries that are appropriate responses to the question.'queries = [ "Delete the record with ID 4 from the 'Staff' table.", 'Delete all records in the "Livestock" table where age is greater than 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# Tokenize the input textsbatch_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']) # normalize embeddingsembeddings = F.normalize(embeddings, p=2, dim=1)scores = (embeddings[:2] @ embeddings[2:].T) * 100print(scores.tolist()){ "Apps": "Given a code contest problem description, retrieve relevant code that can help solve the problem.", "CosQA": "Given a web search query, retrieve relevant code that can help answer the query.", "Text2SQL": "Given a question in text, retrieve SQL queries that are appropriate responses to the question.", "CSN": "Given a piece of code, retrieve the document string that summarizes the code.", "CSN-CCR": "Given a piece of code segment, retrieve the code segment that is the latter part of the code.", "CodeTrans-DL": "Given a piece of code, retrieve code that is semantically equivalent to the input code.", "CodeTrans-Contest": "Given a piece of Python code, retrieve C++ code that is semantically equivalent to the input code.", "StackOverFlow-QA": "Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.", "CodeFeedBack-ST": "Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.", "CodeFeedBack-MT": "Given a multi-turn conversation history that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.", "HummanEval": "Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.", "MBPP": "Given a textual explanation of code functionality, retrieve the corresponding code implementation.", "DS-1000": "Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.", "ODEX": "Given a question, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.", "RepoEval": "Given a piece of code segment, retrieve the code segment that is the latter part of the code.", "SWE-bench-Lite": "Given a code snippet containing a bug and a natural language description of the bug or error, retrieve code snippets that demonstrate solutions or fixes for similar bugs or errors (the desired documents)."}@misc{bge_code, title={Towards A Generalist Code Embedding Model Based On Massive Data Synthesis}, author={Chaofan Li and Jianlyu Chen and Yingxia Shao and Defu Lian and Zheng Liu}, year={2025}, eprint={2505.12697}, archivePrefix={arXiv}, primaryClass={cs.IR}, url={https://arxiv.org/abs/2505.12697}, }