nebula2025/CodeR-full

1---2language:3- zh4- en5tags:6- sentence-transformers7- sentence-similarity8- feature-extraction9- transformers10pipeline_tag: sentence-similarity11library_name: sentence-transformers12license: apache-2.013---14 15Here is the CodeR model trained on both text-only data and the full code data.16 17## Usage18 19### Using FlagEmbedding20 21```22git clone https://github.com/FlagOpen/FlagEmbedding.git23cd FlagEmbedding24pip install -e .25```26 27```python28from FlagEmbedding import FlagLLMModel29queries = [30    "Delete the record with ID 4 from the 'Staff' table.", 31    'Delete all records in the "Livestock" table where age is greater than 5'32]33documents = [34    "DELETE FROM Staff WHERE StaffID = 4;",35    "DELETE FROM Livestock WHERE age > 5;"36]37model = FlagLLMModel('nebula2025/CodeR-full', 38                     query_instruction_format="<instruct>{}\n<query>{}",39                     query_instruction_for_retrieval="Given a question in text, retrieve SQL queries that are appropriate responses to the question.",40                     trust_remote_code=True,41                     use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation42embeddings_1 = model.encode_queries(queries)43embeddings_2 = model.encode_corpus(documents)44similarity = embeddings_1 @ embeddings_2.T45print(similarity)46```47 48By 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.49 50### Using Sentence Transformers51 52```python53from sentence_transformers import SentenceTransformer54import torch55 56# Load the model, optionally in float16 precision for faster inference57model = SentenceTransformer("nebula2025/CodeR-full", model_kwargs={"torch_dtype": torch.float16, "trust_remote_code": True}, tokenizer_kwargs={"trust_remote_code": True})58 59# Prepare a prompt given an instruction60instruction = 'Given a question in text, retrieve SQL queries that are appropriate responses to the question.'61prompt = f'<instruct>{instruction}\n<query>'62# Prepare queries and documents63queries = [64    "Delete the record with ID 4 from the 'Staff' table.", 65    'Delete all records in the "Livestock" table where age is greater than 5'66]67documents = [68    "DELETE FROM Staff WHERE StaffID = 4;",69    "DELETE FROM Livestock WHERE age > 5;"70]71 72# Compute the query and document embeddings73query_embeddings = model.encode(queries, prompt=prompt)74document_embeddings = model.encode(documents)75 76# Compute the cosine similarity between the query and document embeddings77similarities = model.similarity(query_embeddings, document_embeddings)78print(similarities)79```80 81### Using HuggingFace Transformers82 83```python84import torch85import torch.nn.functional as F86 87from torch import Tensor88from transformers import AutoTokenizer, AutoModel89 90 91def last_token_pool(last_hidden_states: Tensor,92                 attention_mask: Tensor) -> Tensor:93    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])94    if left_padding:95        return last_hidden_states[:, -1]96    else:97        sequence_lengths = attention_mask.sum(dim=1) - 198        batch_size = last_hidden_states.shape[0]99        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]100 101 102def get_detailed_instruct(task_description: str, query: str) -> str:103    return f'<instruct>{task_description}\n<query>{query}'104 105 106instruction = 'Given a question in text, retrieve SQL queries that are appropriate responses to the question.'107queries = [108    "Delete the record with ID 4 from the 'Staff' table.", 109    'Delete all records in the "Livestock" table where age is greater than 5'110]111documents = [112    "DELETE FROM Staff WHERE StaffID = 4;",113    "DELETE FROM Livestock WHERE age > 5;"114]115input_texts = queries + documents116 117tokenizer = AutoTokenizer.from_pretrained('nebula2025/CodeR-full', trust_remote_code=True)118model = AutoModel.from_pretrained('nebula2025/CodeR-full', trust_remote_code=True)119model.eval()120 121max_length = 4096122# Tokenize the input texts123batch_dict = tokenizer(input_texts, max_length=max_length, padding=True, truncation=True, return_tensors='pt', pad_to_multiple_of=8)124 125with torch.no_grad():126    outputs = model(**batch_dict)127    embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])128    129# normalize embeddings130embeddings = F.normalize(embeddings, p=2, dim=1)131scores = (embeddings[:2] @ embeddings[2:].T) * 100132print(scores.tolist())133```