Update README.md

README.md

@@ -1,3 +1,86 @@
----
-license: mit
----
+---
+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).