Embedding & Reranking Models
Relevant source files
The following files were used as context for generating this wiki page:
Introduction
The "Embedding & Reranking Models" module provides the semantic foundation for the Agentic GraphRAG system. It is responsible for transforming clinical text, graph entities, and relationships into high-dimensional vector representations and subsequently refining search results through sophisticated reranking strategies. These models enable the system to perform complementary semantic search across four distinct vector collections: entities, relations, chunks, and community reports.
The module utilizes local models via the sentence-transformers library, specifically leveraging the ZeroEntropy suite (Zembed and Zerank) for high-performance medical domain embeddings and reranking. This infrastructure supports multi-strategy retrieval, where atomic search methods are fused and then re-evaluated by cross-encoders to ensure clinical accuracy and relevance.
Sources: README.md:27-33, README.md:121-125, am_diag/vector/reranking/sentence_transformers.py:27-45
Embedding Architecture
The system employs a multi-vector embedding strategy to capture the nuances of a medical knowledge graph. Rather than a single vector, information is vectorized into four specialized collections to support different retrieval granularities.
Vector Collections and Representations
| Collection | Text Embedded | Purpose |
|---|---|---|
| Entity | Canonical name | Semantic lookup of medical concepts (Drugs, Diseases, etc.) |
| Relation | Edge fact, Edge type, and Full SPO | Captures predicates and participant relationships |
| Chunk | Raw chunk text | Traditional dense retrieval over document passages |
| CommunityReport | Report summary | Thematic, cluster-level retrieval for broad questions |
Sources: README.md:121-137, am_diag/ingestion/corpus_ingestion.py:38-89
Zembed Specialization
The ZembedEmbedder extends standard sentence-transformers capabilities with support for MATRYOSHKA-style embeddings. It allows for dynamic dimensionality reduction (from 2560 down to 80) using a learned projection matrix and supports multiple quantization precisions (int8, uint8, binary).
This diagram illustrates the ZembedEmbedder logic, where vectors can be projected to lower dimensions and quantized to save storage while maintaining semantic integrity.
Sources: am_diag/vector/embedding/zembed.py, tests/am_diag/vector/embedding/test_zembed.py:46-120
Reranking Pipeline
Reranking serves as the second stage of retrieval, where a Cross-Encoder evaluates the specific relevance of a query against a set of candidate documents. Unlike bi-encoder embeddings, the cross-encoder processes the query and document simultaneously to produce a more accurate similarity score.
SentenceTransformersReranker
The SentenceTransformersReranker class manages local CrossEncoder models (defaulting to zeroentropy/zerank-2-reranker). Key features include:
- Instruction Following: Supports natural-language instructions (e.g., "Prioritize medical accuracy") via
<instruction>tags. - Lazy Loading: The model is initialized only on the first
rerankcall and cached. - Thread Safety: Uses
asyncio.to_threadto execute synchronous GPU/CPU-bound inference without blocking the async event loop. - Calibration: Raw logits from the model are calibrated using a sigmoid function (scaled by 5.0) to produce a 0.0-1.0 relevance score.
Sources: am_diag/vector/reranking/sentence_transformers.py:27-145, tests/am_diag/vector/reranking/test_sentence_transformers.py:105-135
Reranking Data Flow
The sequence shows how instructions are injected into the query before the cross-encoder computes calibrated scores for document ranking.
Sources: am_diag/vector/reranking/sentence_transformers.py:97-152, tests/am_diag/vector/reranking/test_sentence_transformers.py:120-145
Implementation Details
Model Configuration
Both embedding and reranking components are configured via pydantic_settings.
| Parameter | Default Value | Description |
|---|---|---|
model_name | zeroentropy/zerank-2-reranker | The HuggingFace model path |
device | cpu | Hardware for inference (cpu, cuda) |
batch_size | 32 | Pairs processed per internal model call |
max_length | 32768 | Maximum token context window |
trust_remote_code | False | Whether to allow custom HF code |
Sources: am_diag/vector/reranking/sentence_transformers.py:48-60, am_diag/vector/embedding/zembed.py:20-25
Data Models
The module interacts with standardized data models defined in the core package:
- VectorRecord: Contains the
id,vector(list of floats), andpayload(metadata) for upsertion into stores like Qdrant or Weaviate. - RerankResult: Encapsulates the
index,score, andtextof a document after reranking.
Sources: am_diag/common/data_models/init.py:75-76, am_diag/vector/reranking/base.py
Integration with Ingestion
During the ingestion process, the _embed_and_upsert function coordinates the vectorization of extracted knowledge. It ensures that every entity, relation, and text chunk produced by the extraction pipeline is indexed for future search.
# Simplified flow from am_diag/ingestion/corpus_ingestion.py
entity_vecs = await embedder.embed(entity_texts)
entity_records = [
VectorRecord(id=ent.id, vector=vec, payload=pl)
for ent, vec, pl in zip(extracted_entities, entity_vecs, entity_payloads)
]
await vector_store.upsert(rconfig.entity_collection, entity_records)
Sources: am_diag/ingestion/corpus_ingestion.py:75-84
Summary
The Embedding & Reranking Models provide a robust semantic layer for the medical diagnosis system. By utilizing specialized medical models and multi-faceted representations of the knowledge graph, the system ensures that retrieval is not only efficient but also clinically relevant. The architecture's support for MATRYOSHKA embeddings and instruction-aware reranking allows it to balance performance with precision across varying scales of medical literature.