cs4248-nlp/paper-s3-a3-bimga-query-only-tinybert-general-4l-312d-taco-hf-20260402-015143
016
cs4248-nlp/paper-s3-a3-bimga-query-only-tinybert-general-4l-312d-taco-hf-20260402-015143
Code-search embedding model trained with the CS4248 two-phase KD pipeline.
Model details
Usage
from transformers import AutoModel, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("cs4248-nlp/paper-s3-a3-bimga-query-only-tinybert-general-4l-312d-taco-hf-20260402-015143")
model = AutoModel.from_pretrained("cs4248-nlp/paper-s3-a3-bimga-query-only-tinybert-general-4l-312d-taco-hf-20260402-015143")Mean-pool the last hidden state to get a fixed-size embedding:
import torch
def mean_pool(model_output, attention_mask):
token_embeddings = model_output.last_hidden_state
mask = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
return (token_embeddings * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
inputs = tokenizer("your query here", return_tensors="pt", truncation=True, max_length=160)
with torch.no_grad():
outputs = model(**inputs)
embedding = mean_pool(outputs, inputs['attention_mask'])