32.37M parameter text embedding model by
Axiotic AI, achieving
57.38 average on MTEB English (66/66 tasks).
9-layer transformer, 512 hidden dim, mean pooling — strongest overall model.
1import torch
2from huggingface_hub import snapshot_download
3import sys, yaml
4
5# Download model from HuggingFace
6model_path = snapshot_download("axiotic/ogma-large")
7sys.path.insert(0, model_path)
8
9from ogma_model import OgmaModel
10from config import OgmaConfig, TaskToken
11from tokenizer import OgmaTokenizer
12
13# Load model
14with open(f"{model_path}/config.yaml") as f:
15 cfg = yaml.safe_load(f)
16config = OgmaConfig.from_dict(cfg)
17model = OgmaModel(config)
18state = torch.load(f"{model_path}/model.pt", map_location="cpu", weights_only=True)
19model.load_state_dict(state)
20model.eval()
21
22# Load tokenizer
23tokenizer = OgmaTokenizer(f"{model_path}/tokenizer.json")
24
25# Encode text
26sentences = ["The quick brown fox", "A fast auburn canine"]
27enc = tokenizer.batch_encode(sentences, max_length=1024)
28ids = torch.tensor(enc["input_ids"])
29mask = torch.tensor(enc["attention_mask"])
30
31with torch.no_grad():
32 embs = model.encode(ids, mask, task=TaskToken.SYM)
33
34# Cosine similarity
35sim = torch.nn.functional.cosine_similarity(embs[0], embs[1], dim=0)
36print(f"Similarity: {sim.item():.4f}")
37print(f"Shape: {embs.shape}") # (2, 256)
1queries = ["What is machine learning?"]
2documents = ["ML is a subset of AI...", "The weather is sunny today"]
3
4q_enc = tokenizer.batch_encode(queries, max_length=1024)
5d_enc = tokenizer.batch_encode(documents, max_length=1024)
6
7with torch.no_grad():
8 # Symmetric: both queries and documents use TaskToken.QRY (not a typo).
9 # Swap TaskToken.QRY → TaskToken.SYM on both sides to try the SYM route instead.
10 q_embs = model.encode(torch.tensor(q_enc["input_ids"]),
11 torch.tensor(q_enc["attention_mask"]), task=TaskToken.QRY)
12 d_embs = model.encode(torch.tensor(d_enc["input_ids"]),
13 torch.tensor(d_enc["attention_mask"]), task=TaskToken.QRY)
14
15scores = q_embs @ d_embs.T
16print(f"Relevance scores: {scores}")
1full = model.encode(ids, mask, task=TaskToken.SYM) # (256d)
2small = torch.nn.functional.normalize(full[:, :32], p=2, dim=-1) # (32d)
Benchmarked with MTEB v2.10.7 on the standard 66-task English benchmark using category averaging (same methodology as the MTEB leaderboard).
This model is licensed under
CC-BY-NC-4.0. Commercial use requires a separate license from Axiotic AI.