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.7.tar.gz (8.6 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.7-py3-none-any.whl (8.0 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: mizan_embedder-0.2.7.tar.gz
  • Upload date:
  • Size: 8.6 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.7.tar.gz
Algorithm Hash digest
SHA256 5cbb4fe631640c68a8cffbebab8c5b6b3c4c50a28024874038300f9ff2a0c08d
MD5 5dc43f49adf37469efe3670ea6022666
BLAKE2b-256 c5e1750d9a5c3e41f3267db3adf813a9bc02d5343513bbd9ef7118bd0f0d2d9a

See more details on using hashes here.

File details

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

File metadata

  • Download URL: mizan_embedder-0.2.7-py3-none-any.whl
  • Upload date:
  • Size: 8.0 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.7-py3-none-any.whl
Algorithm Hash digest
SHA256 4c02db6dccec2f92454da1e2d29ad4476cdd655544a2678ae2ad48956b122405
MD5 25985d5f5bb66024fd1783dbd4c6d144
BLAKE2b-256 19ee3918bbc0a976561183285b26de7e4bca75033828d0720ea1dea9bff96756

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