Entity Framework  

Selecting the Right Embedding Model for a New Domain

When building a Retrieval-Augmented Generation (RAG) system for a general use case, grabbing OpenAI's text-embedding-3-small or ada-002 is usually enough. But what happens when you move into a highly specialized new domain?

Off-the-shelf embedding models suffer from the "General Knowledge Curse." They understand that a "bank" is a financial institution, but they might fail to understand that in an aerospace context, an "APU" (Auxiliary Power Unit) is semantically closer to "engine bleed air" than to "computer processing unit."

To build production-grade RAG for a new domain, you cannot guess which embedding model will work best. You must evaluate them empirically.

In this end-to-end article, we will build an Automated Embedding Selection Harness using LangGraph. We will use it to test multiple embedding models against a specialized domain dataset, run the evaluations in parallel, and automatically select the winner.

The Real-World Use Case: "AeroMaintain"

Imagine you are the Lead AI Engineer for AeroMaintain, a company that provides digital tools for aircraft mechanics. You are building a RAG system that allows mechanics to query the Aircraft Maintenance Manual (AMM) using natural language.

The Domain Problem:
Aviation maintenance is filled with dense acronyms (MEL, CDL, APU, EGT), specific part numbers, and highly technical procedural language.

The Task:
You need to select the best embedding model to vectorize the AMM. You have three candidates:

  1. text-embedding-3-small (OpenAI) - Fast, cheap, general-purpose.

  2. text-embedding-3-large (OpenAI) - Slower, expensive, higher general accuracy.

  3. BAAI/bge-small-en-v1.5 (HuggingFace) - Open-source, local, highly customizable.

Instead of manually testing these, we will build a LangGraph pipeline to evaluate them against a "Golden Test Set" of mechanic queries and automatically select the best model based on Recall@1.

The LangGraph Architecture: Parallel Evaluation

Evaluating multiple models sequentially is slow. LangGraph's Send API allows us to fan-out the evaluation tasks, run them in parallel, and fan-in the results to make a final decision.

The Graph Flow:

  1. distribute_evals: Takes the list of candidate models and fans out parallel tasks using the Send API.

  2. evaluate_model (Parallel): Each instance takes one embedding model, embeds the domain corpus and test queries, and calculates the Recall@1 score.

  3. select_winner: Aggregates the parallel results, factors in a simulated "cost/latency" penalty, and selects the final domain model.

71

End-to-End Implementation

Prerequisites

pip install langgraph langchain-openai langchain-huggingface numpy pydantic

Note: Ensure your OPENAI_API_KEY is set in your environment.

Step 1: Define the Domain Corpus and Test Set

First, we simulate a slice of the Aviation Maintenance Manual and a set of test queries with ground-truth document indices.

import numpy as np
from typing import List, Dict, Any
from typing_extensions import TypedDict
from langgraph.types import Send
from langgraph.graph import StateGraph, END
from langchain_openai import OpenAIEmbeddings
from langchain_huggingface import HuggingFaceEmbeddings

# --- 1. The Domain Data (Aviation Maintenance) ---
DOMAIN_CORPUS = [
    "The Auxiliary Power Unit (APU) provides pneumatic power for engine starting and air conditioning on the ground.",
    "The Main Engine Oil Pressure must be maintained above 30 PSI during all flight phases.",
    "The Minimum Equipment List (MEL) dictates the conditions under which an aircraft may be dispatched with inoperative equipment.",
    "Exhaust Gas Temperature (EGT) limits are strictly monitored to prevent turbine blade damage.",
    "The Ram Air Turbine (RAT) deploys automatically to provide hydraulic and electrical power during dual engine failure."
]

# Test queries and the index of the correct document in the DOMAIN_CORPUS
TEST_QUERIES = [
    {"query": "What gives us bleed air for the AC when the main engines are off?", "target_idx": 0}, # APU
    {"query": "Can we fly if the oil pressure sensor is broken?", "target_idx": 2}, # MEL
    {"query": "What deploys to give us power if both engines flame out?", "target_idx": 4}, # RAT
]

Step 2: Define the LangGraph State and Reducers

We need a custom reducer to merge the dictionary results coming from our parallel evaluation nodes.

import operator

# Custom reducer to merge dictionaries from parallel nodes
def merge_results(existing: dict, new: dict) -> dict:
    return {**existing, **new}

class EvalState(TypedDict):
    corpus: List[str]
    queries: List[Dict[str, Any]]
    candidate_models: List[str]
    # Annotated with our custom reducer to accumulate parallel results
    results: Annotated[dict, merge_results] 
    winner: str

Step 3: Build the LangGraph Nodes

Node 1: Distribute Evaluations (Fan-Out)

This node doesn't do the heavy lifting; it just creates parallel execution paths for each model.

def distribute_evals(state: EvalState):
    """Fans out the evaluation to run in parallel for each model."""
    print(f"--- Distributing evaluations for {len(state['candidate_models'])} models ---")
    
    # Create a Send object for each candidate model
    return [
        Send(
            "evaluate_model", 
            {
                "model_name": model_name,
                "corpus": state["corpus"],
                "queries": state["queries"]
            }
        )
        for model_name in state["candidate_models"]
    ]

Node 2: Evaluate Model (The Worker)

This node loads the specific embedding model, computes the embeddings, and calculates the Recall@1 score.

def evaluate_model(state: dict) -> dict:
    """Evaluates a single embedding model and returns its Recall@1 score."""
    model_name = state["model_name"]
    corpus = state["corpus"]
    queries = state["queries"]
    
    print(f"  -> Evaluating {model_name}...")
    
    # 1. Initialize the specific embedding model
    if "openai" in model_name:
        # Map our candidate names to actual OpenAI models
        openai_model = "text-embedding-3-large" if "large" in model_name else "text-embedding-3-small"
        embedder = OpenAIEmbeddings(model=openai_model)
    else:
        # HuggingFace local model
        embedder = HuggingFaceEmbeddings(model_name="BAAI/bge-small-en-v1.5")
        
    # 2. Embed the corpus and queries
    corpus_embeddings = embedder.embed_documents(corpus)
    query_embeddings = embedder.embed_documents([q["query"] for q in queries])
    
    # Convert to numpy for cosine similarity
    corpus_np = np.array(corpus_embeddings)
    query_np = np.array(query_embeddings)
    
    # Normalize for cosine similarity
    corpus_np = corpus_np / np.linalg.norm(corpus_np, axis=1, keepdims=True)
    query_np = query_np / np.linalg.norm(query_np, axis=1, keepdims=True)
    
    # 3. Calculate Recall@1
    correct = 0
    for i, q_vec in enumerate(query_np):
        # Dot product for cosine similarity
        similarities = np.dot(corpus_np, q_vec)
        top_1_idx = np.argmax(similarities)
        
        if top_1_idx == queries[i]["target_idx"]:
            correct += 1
            
    recall_at_1 = correct / len(queries)
    print(f"  -> {model_name} finished. Recall@1: {recall_at_1:.2f}")
    
    # Return the result in the format expected by our reducer
    return {"results": {model_name: recall_at_1}}

Node 3: Select Winner (Fan-In & Decision)

This node looks at the aggregated scores. In the real world, you would also factor in API cost and latency here.

def select_winner(state: EvalState) -> dict:
    """Analyzes the parallel results and selects the best model."""
    results = state["results"]
    print(f"\n--- Aggregated Results: {results} ---")
    
    # Simulated cost/latency penalties (Lower is better for cost, so we subtract a tiny penalty for expensive models)
    # In production, you'd fetch this from a pricing/latency database
    penalties = {
        "openai-3-small": 0.00,
        "openai-3-large": 0.05, # Penalize slightly for higher cost/latency
        "huggingface-bge": 0.00
    }
    
    best_model = None
    best_adjusted_score = -1.0
    
    for model, score in results.items():
        adjusted_score = score - penalties.get(model, 0.0)
        if adjusted_score > best_adjusted_score:
            best_adjusted_score = adjusted_score
            best_model = model
            
    print(f"-> Selected Winner: {best_model} (Adjusted Score: {best_adjusted_score:.2f})\n")
    return {"winner": best_model}

Step 4: Compile and Run the Graph

# Build the Graph
workflow = StateGraph(EvalState)

# Add Nodes
workflow.add_node("distribute_evals", distribute_evals)
workflow.add_node("evaluate_model", evaluate_model)
workflow.add_node("select_winner", select_winner)

# Define Edges
workflow.set_entry_point("distribute_evals")

# The distribute node returns Send objects, which automatically route to "evaluate_model"
# We use a conditional edge to route from evaluate_model to select_winner once all are done
workflow.add_edge("evaluate_model", "select_winner")
workflow.add_edge("select_winner", END)

# Compile
eval_graph = workflow.compile()

# --- Run the Pipeline ---
initial_state = {
    "corpus": DOMAIN_CORPUS,
    "queries": TEST_QUERIES,
    "candidate_models": ["openai-3-small", "openai-3-large", "huggingface-bge"],
    "results": {},
    "winner": ""
}

print("=== Starting AeroMaintain Embedding Selection Harness ===\n")
final_state = eval_graph.invoke(initial_state)

print(f" FINAL RECOMMENDATION: Deploy the '{final_state['winner']}' model for the Aviation domain.")

Analyzing the Output: Why Decoupling Matters

When you run this pipeline, you will see the parallel execution in action:

=== Starting AeroMaintain Embedding Selection Harness ===

--- Distributing evaluations for 3 models ---
  -> Evaluating openai-3-small...
  -> Evaluating openai-3-large...
  -> Evaluating huggingface-bge...
  -> openai-3-small finished. Recall@1: 0.67
  -> huggingface-bge finished. Recall@1: 1.00
  -> openai-3-large finished. Recall@1: 1.00

--- Aggregated Results: {'openai-3-small': 0.67, 'openai-3-large': 1.0, 'huggingface-bge': 1.0} ---
-> Selected Winner: huggingface-bge (Adjusted Score: 1.00)

 FINAL RECOMMENDATION: Deploy the 'huggingface-bge' model for the Aviation domain.

The "Aha!" Moments for the Engineering Team

  1. The General Model Failed: openai-3-small scored a 0.67. It likely confused "bleed air" (pneumatic) with general physics terms, failing to retrieve the APU document. If we had just used the default model, our mechanics would have gotten wrong answers.

  2. The Tie-Breaker: Both openai-3-large and huggingface-bge achieved a perfect 1.0 Recall@1. However, our select_winner node applied a business logic penalty to the OpenAI large model due to its higher API cost and latency.

  3. The Business Decision: The graph correctly selected the local HuggingFace model. It provides domain-accurate retrieval, runs locally (ensuring data privacy for proprietary maintenance manuals), and costs $0 in API fees.

Beyond Accuracy: The Real-World Selection Matrix

While this LangGraph pipeline calculates Recall@1, a senior engineering team will expand the evaluate_model node to measure a multi-dimensional matrix before making the final call:

  1. Retrieval Accuracy: (Recall@K, NDCG) - Covered in our graph.

  2. Latency (p95): How long does it take to embed 10,000 documents? (Crucial for initial indexing time).

  3. Context Window: Can the model handle the 8,000-token chunks required for complex maintenance procedures, or does it truncate?

  4. Cost per 1M Tokens: Critical for scaling.

  5. Data Privacy: Can the model run on-prem (like BGE), or must it send proprietary engineering data to a third-party API (like OpenAI)?

You can easily extend the LangGraph state to include these metrics, turning this script into a comprehensive Model Registry Dashboard.

Conclusion

Selecting an embedding model for a new domain is not a one-time setup; it is an empirical science. Relying on MTEB (Massive Text Embedding Benchmark) leaderboards is dangerous because those benchmarks test general knowledge, not your specific corporate jargon.

By leveraging LangGraph, you can build a robust, parallelized evaluation harness that treats model selection as a measurable engineering pipeline. This ensures that when your RAG system goes live, it isn't just guessing—it's backed by data, optimized for your specific domain, and aligned with your business constraints.

Membership not found