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 MizanVector 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:

  • mizanvector (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 mizanvector.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 mizanvector
Example: full semantic search pipeline

from mizanvector 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.0.tar.gz (6.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.0-py3-none-any.whl (6.9 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: mizan_embedder-0.2.0.tar.gz
  • Upload date:
  • Size: 6.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.0.tar.gz
Algorithm Hash digest
SHA256 be75d1b6daff5b1cca363eaa9ca1edf16b56f32762353fb0c3d521aa54b27e29
MD5 870b0b9f6fe171d55f53b4d3fded598b
BLAKE2b-256 d09dc53117c0b635d20948c6a17ad0b993dedc850ca9e87a198459dd4eb6aad2

See more details on using hashes here.

File details

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

File metadata

  • Download URL: mizan_embedder-0.2.0-py3-none-any.whl
  • Upload date:
  • Size: 6.9 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.0-py3-none-any.whl
Algorithm Hash digest
SHA256 1a9f75dc7d0f50d9a1a717d4e3682f8b04f616b105a90397207b9f884b05f56a
MD5 92bdea53cbd2400bbeee264a4deceb79
BLAKE2b-256 bce95361e49e54a7f484c3c9c9f8e91f8a4f3f5038a6a67a93eb2664c2d4d656

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