SageLite-s

SageLite-s

Model Description

SageLite is a new family of open embedding models with an encoder architecture that supports a wide range of tasks in both code and text. SageLite went through three stages of training:

1. MLM Pretraining: Standard masked language model (MLM) pretraining on mixed code and text data (The-Stack-v2 and Falcon-refinedweb).
2. Contrastive Pre-Finetuning: Learning from a large amount of positive pairs mined from web data and GitHub.
3. Contrastive Fine-Tuning: Fine-tuning on a small amount of synthetic data.

---

Training Data

This checkpoint is trained on both The-Stack-v2 and Falcon-refinedweb. Supported languages (15 in total) are: English, C, C#, Go, Java, JavaScript, TypeScript, PHP, Python, and Ruby.

---

How to Use

This checkpoint consists of an encoder (80M model) that extracts code embeddings of 768 dimensions. It can be loaded using the Hugging Face Transformers library and employs the Starcoder Tokenizer.

from transformers import AutoModel, AutoTokenizer

# Specify the checkpoint
checkpoint = "SageLite/SageLite-s"
device = "cuda"  # Use "cpu" if GPU is unavailable

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True, add_eos_token=True)
model = AutoModel.from_pretrained(checkpoint, trust_remote_code=True).to(device)

# Example usage
code_snippet = "def print_hello_world():\tprint('Hello World!')"
inputs = tokenizer.encode(code_snippet, return_tensors="pt").to(device)
embedding = model(inputs)[0]  # Extract the embedding

Code Retrieval Performance

1. Code2Code Search

2. NL2Code Search

---

Text Retrieval Performance (MTEB Retrieval)

---