| --- |
| license: mit |
| --- |
| # Cosine-Embed |
|
|
| Cosine-Embed is a PyTorch sentence embedding model trained to place similar texts close together in an embedding space. The model outputs L2-normalized vectors so cosine similarity is computed as a dot product. |
|
|
| ## What it produces |
| - Input: tokenized text (`input_ids`, `attention_mask`) |
| - Output: an embedding vector of size `hidden_dim` with L2 normalization |
| - Cosine similarity: `cos(a, b) = embedding(a) · embedding(b)` |
|
|
| ## Model details |
| - Transformer blocks (custom implementation using RMSNorm, RoPE positional encoding, and SwiGLU feed-forward) |
| - Masked mean pooling over token embeddings |
| - Final L2 normalization |
|
|
| ## Default configuration |
| These parameters are used in `Notebooks/Training.ipynb`: |
| - `vocab_size`: 30522 |
| - `seq_len`: 128 |
| - `hidden_dim`: 512 |
| - `n_heads`: 8 |
| - `n_layer`: 3 |
| - `ff_dim`: 2048 |
| - `eps`: 1e-5 |
| - `dropout`: 0.1 |
|
|
| ## Training objective |
| The model is trained with triplet loss on cosine similarity: |
|
|
| `loss = max(0, sim(anchor, negative) - sim(anchor, positive) + margin)` |
|
|
| ## Checkpoints |
| - `checkpoints/checkpoint.pt`: training checkpoint (model, optimizer, losses, and configs) |
| - `checkpoints/model.safetensors`: weights-only export for inference |
|
|
| ## Minimal inference |
| ```python |
| import torch |
| from transformers import AutoTokenizer |
| from safetensors.torch import load_file |
| |
| from Architecture import EmbeddingModel, ModelConfig |
| |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| |
| state_dict = load_file("checkpoints/model.safetensors") |
| |
| cfg = ModelConfig( |
| vocab_size=30522, |
| seq_len=128, |
| hidden_dim=512, |
| n_heads=8, |
| n_layer=3, |
| eps=1e-5, |
| ff_dim=2048, |
| dropout=0.1, |
| ) |
| |
| model = EmbeddingModel(cfg).to(device) |
| model.load_state_dict(state_dict) |
| model.eval() |
| |
| tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") |
| |
| def embed(texts): |
| enc = tokenizer( |
| texts, |
| padding=True, |
| truncation=True, |
| max_length=128, |
| return_tensors="pt", |
| ) |
| enc = {k: v.to(device) for k, v in enc.items()} |
| with torch.no_grad(): |
| return model(enc["input_ids"], enc["attention_mask"]) # normalized |
| |
| def cosine_similarity(a, b): |
| ea = embed([a])[0] |
| eb = embed([b])[0] |
| return float((ea * eb).sum().item()) |
| ``` |
|
|
| ## Notes |
| - Use the same tokenizer (`bert-base-uncased`) and the same `max_length=128` (or keep `seq_len` and preprocessing consistent). |
|
|
