--- license: cc-by-nc-sa-4.0 tags: - sentence-transformers - transformers - splade - sparse-encoder - code pipeline_tag: feature-extraction --- SPLADE-Code-8B is a sparse retrieval model designed for code retrieval tasks. ## Usage ### Using Sentence Transformers Install Sentence Transformers: ```bash pip install sentence_transformers ``` ```python from sentence_transformers import SparseEncoder model = SparseEncoder("naver/splade-code-8B", trust_remote_code=True) queries = [ "SELECT *\nFROM Student\nWHERE Age = (\nSELECT MAX(Age)\nFROM Student\nWHERE Group = 'specific_group'\n)\nAND Group = 'specific_group';" ] query_embeddings = model.encode(queries) print(query_embeddings.shape) # torch.Size([1, 151936]) sparsity = model.sparsity(query_embeddings) print(sparsity) # {'active_dims': 1120.0, 'sparsity_ratio': 0.9926284751474305} decoded = model.decode(query_embeddings, top_k=10) print(decoded) # [[ # ('Ġgroup', 2.34375), # ('Ġoldest', 2.28125), # ('Ġage', 2.25), # ('_group', 2.25), # ('ĠGroup', 2.171875), # ('ĠAge', 2.109375), # ('ĠMAX', 2.0625), # ('ĠStudent', 2.046875), # ('Ġspecific', 2.03125), # ('Ġstudent', 2.0), # ]] ``` ### Using Transformers ```bash pip install transformers ``` ```python from transformers import AutoModelForCausalLM, AutoModel import os import torch splade = AutoModelForCausalLM.from_pretrained("naver/splade-code-8B", trust_remote_code=True) device = (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")) splade.to(device) splade.eval() queries = ["SELECT *\nFROM Student\nWHERE Age = (\nSELECT MAX(Age)\nFROM Student\nWHERE Group = 'specific_group'\n)\nAND Group = 'specific_group';"] bow_dict = splade.encode(queries, prompt_type="query", top_k_q=10, return_dict=True, print_dict=True) ``` ``` +--------------------------------------------------------------------+ | TOP ACTIVATED WORDS | +--------------------------------------------------------------------+ * INPUT: SELECT * FROM Student WHERE Age = ( SELECT MAX(Age) FROM Student WHERE Group = 'specific_group' ) AND Group = 'specific_group'; Ġgroup | ████████████████████ 2.34 Ġoldest | ███████████████████ 2.28 Ġage | ███████████████████ 2.25 _group | ███████████████████ 2.25 ĠGroup | ██████████████████ 2.17 ĠAge | ██████████████████ 2.11 ĠMAX | █████████████████ 2.06 ĠStudent | █████████████████ 2.05 Ġspecific | █████████████████ 2.03 Ġstudent | █████████████████ 2.00 ```