naver/splade-code-8B

SPLADE-Code-8B is a sparse retrieval model designed for code retrieval tasks.

Usage

Using Sentence Transformers

Install Sentence Transformers:

pip install sentence_transformers

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

pip install transformers

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