Skip to main content

Mizan-optimized embedding model training

Project description

mizan-embedder

Mizan-optimized Embedding Models for AI, Search, and RAG.
mizan-embedder is the official embedding-model library in the Mizan ecosystem, designed to create scale-aware, noise-resistant, and proportionally accurate embeddings trained using the Mizan Balance Function.

Proposed & Developed By:
Ahsan Shaokat โ€” Computer Scientist & AI/ML Researcher
Inventor of the Mizan Balance Function (2025)


๐ŸŒŸ Overview

Modern embedding systems (MiniLM, MPNet, E5, etc.) use cosine similarity, which:

  • โŒ Ignores magnitude
  • โŒ Fails with noisy or multi-scale embeddings
  • โŒ Produces unstable rankings in RAG
  • โŒ Forces L2-normalization (losing information)

Mizan-Embedder fixes this by training models specifically for:

  • โœ” Mizan similarity (scale-aware)
  • โœ” Proportional contrastive learning
  • โœ” Chunk-length stable retrieval
  • โœ” Large document embeddings
  • โœ” Multimodal (text + images)

This library enables you to build your own embedding models, optimized for the mizan_vector search engine.


๐Ÿ“ฆ Features

๐Ÿง  MizanEmbeddingModel-v1

  • Transformer backbone (DistilBERT, MiniLM, BERT, or any HF model)
  • Projection head to target embedding dimension (e.g., 384)
  • Supports mean, cls, and max pooling
  • Optional L2 normalization (usually disabled for Mizan)

๐Ÿงฐ Utilities Included

  • Dataset utilities for contrastive text pairs
  • Collate function for fast tokenization
  • Inference wrapper (MizanTextEncoderWrapper)
  • Example training script (train_text_contrastive.py)

๐Ÿ”Œ Integrates Seamlessly With:

  • mizan_vector (Memory store + Postgres pgvector)
  • mizan-rag (retrieval pipelines)
  • Any Python ML workflow

๐Ÿ“ Project Structure

mizan-embedder/ โ”‚ โ”œโ”€โ”€ mizan_embedder/ โ”‚ โ”œโ”€โ”€ init.py โ”‚ โ”œโ”€โ”€ model.py # MizanEmbeddingModel + inference wrapper โ”‚ โ”œโ”€โ”€ data.py # Dataset + collate functions โ”‚ โ”œโ”€โ”€ train_text_contrastive.py # Example training script โ”œโ”€โ”€ pyproject.toml # PyPI-ready config โ”œโ”€โ”€ README.md โ””โ”€โ”€ LICENSE


โš™๏ธ Installation

From local repo:

pip install -e .
๐Ÿงฑ Architecture
๐Ÿ”น MizanEmbeddingModel
A transformer-based encoder with:

Backbone (HuggingFace model)

Projection layer โ†’ [hidden_size] โ†’ [embedding_dim]

Pooling (mean, cls, max)

Normalization (optional)

Diagram:

mathematica
Input Text โ†’ Tokenizer โ†’ Transformer Backbone โ†’ Pooling โ†’ Projection โ†’ Embedding
๐Ÿ”น Why Projection?
To unify embedding dimensions across:

text models

code models

multimodal models

future Mizan models

๐Ÿš€ Usage
๐Ÿ”น Load the encoder

from mizan_embedder.model import MizanTextEncoderWrapper

encoder = MizanTextEncoderWrapper(
    backbone_name="distilbert-base-uncased",
    emb_dim=384,
    pooling="mean",
    normalize=False,  # Mizan works best without normalization
)

vector = encoder.encode_one("Mizan is a scale-aware similarity function.")
print(vector.shape)
๐Ÿงช Training Your First Mizan Encoder
Use the provided script:

python train_text_contrastive.py
This script:

Loads text pairs

Tokenizes them

Trains with MizanContrastiveLoss

Prints loss per epoch

Example Training Code (simplified)

from mizan_embedder.model import MizanEmbeddingModel
from mizan-vector.losses import MizanContrastiveLoss

model = MizanEmbeddingModel(
    backbone_name="distilbert-base-uncased",
    emb_dim=384,
    pooling="mean",
)

loss_fn = MizanContrastiveLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)

for enc1, enc2, labels in loader:
    emb1 = model(**enc1)
    emb2 = model(**enc2)

    loss = loss_fn(emb1, emb2, labels)
    loss.backward()
    optimizer.step()
๐Ÿ” Contrastive Dataset Format
Your dataset should consist of (text1, text2, label) pairs:

Label = 1 โ†’ similar

Label = 0 โ†’ not similar

Example:

pairs = [
    ("what is mizan?", "mizan is a scale-aware similarity function", 1),
    ("who invented mizan?", "Ahsan Shaokat proposed the Mizan Balance Function", 1),
    ("cosine similarity", "apples are fruit", 0),
]
Dataset loader handles this automatically.

๐Ÿค– Inference: Encoding Many Sentences

texts = [
    "Mizan is scale-aware.",
    "Cosine ignores magnitude.",
    "Apples are fruit.",
]

embs = encoder.encode(texts)
print(embs.shape)  # e.g. torch.Size([3, 384])
๐Ÿ”— Integrating With mizan-vector
Example: full semantic search pipeline

from mizan-vector import MizanMemoryStore
from mizan_embedder.model import MizanTextEncoderWrapper

encoder = MizanTextEncoderWrapper()
store = MizanMemoryStore(dim=384)

docs = [
    "Mizan Balance Function is scale-aware.",
    "Cosine similarity uses only angle.",
    "Ahsan Shaokat invented Mizan.",
]

embs = encoder.encode(docs)

for doc, emb in zip(docs, embs):
    store.add_document(content=doc, embedding=emb.tolist())

query = "who created the mizan function?"
q_emb = encoder.encode_one(query).tolist()

results = store.search(q_emb, top_k=3, metric="mizan")

for r in results:
    print(r.score, r.content)
๐Ÿ”ฅ Why Use Mizan-Based Embeddings?
Problem in Cosine Models	Mizan Solution
Loses magnitude info	Keeps scale meaningfully
Sensitive to noise/outliers	Proportional + stable
Long chunks score lower	Corrects length bias
Normalized embeddings only	No normalization needed
RAG retrieval unstable	Stable across chunk sizes
Cosine โ‰  semantic meaning	Mizan captures proportional similarity

Mizan-optimized embeddings simply behave more naturally for real-world retrieval.

๐Ÿ—บ๏ธ Roadmap
Next Versions:
โœ” MizanTextEncoder-base-384
STS/NLI-trained

Released in mizan-models

โœ” MizanCodeEncoder-base
CodeBERT-based

Code โ†” docstring training

โœ” MizanMultimodalEncoder-v1
CLIP-based

Image โ†” text contrastive training

โœ” mizan-rag
Full retrieval pipeline (chunking โ†’ embedding โ†’ storing โ†’ LLM answering)

๐Ÿ“œ License
MIT License
ยฉ 2025 Ahsan Shaokat

๐Ÿ™Œ Acknowledgements
Special thanks to:

HuggingFace transformers

pgvector open-source community

PyTorch developers

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

mizan_embedder-0.2.9.tar.gz (8.3 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

mizan_embedder-0.2.9-py3-none-any.whl (7.7 kB view details)

Uploaded Python 3

File details

Details for the file mizan_embedder-0.2.9.tar.gz.

File metadata

  • Download URL: mizan_embedder-0.2.9.tar.gz
  • Upload date:
  • Size: 8.3 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.13.3

File hashes

Hashes for mizan_embedder-0.2.9.tar.gz
Algorithm Hash digest
SHA256 ee89719f2fae9913b11692364d1cacd58bc8f9aeb7e194e5d4cbdca86b7d3e3d
MD5 2ab2a07455497f075e37f1628f35bba1
BLAKE2b-256 2d00adaeea866eb593d29109431dee2b6151906c813ccaeabdf8c59c3361f78e

See more details on using hashes here.

File details

Details for the file mizan_embedder-0.2.9-py3-none-any.whl.

File metadata

  • Download URL: mizan_embedder-0.2.9-py3-none-any.whl
  • Upload date:
  • Size: 7.7 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.13.3

File hashes

Hashes for mizan_embedder-0.2.9-py3-none-any.whl
Algorithm Hash digest
SHA256 51b78b2ee00bf35c53373d9a9097adec5b58de82edec67e16d5afde8fd67b16e
MD5 e4c12fa41d33c11e8e6dafdb5a70b9c5
BLAKE2b-256 eef022c5574461339488b25d4de76116edfe276cc4843917a413c475bd0c66b8

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page