Backends

novelentitymatcher.backends.base

Classes

RerankerBackend

Bases: ABC

Abstract base class for reranker backends.

Functions

score(query, docs) abstractmethod

Score query-document pairs.

Source code in src/novelentitymatcher/backends/base.py

@abstractmethod
def score(self, query: str, docs: list[str]) -> list[float]:
    """Score query-document pairs."""

rerank(query, candidates, top_k=5, text_field='text')

Rerank candidates and return top_k.

Default implementation using score(). Subclasses can override for optimization.

Source code in src/novelentitymatcher/backends/base.py

def rerank(
    self,
    query: str,
    candidates: list[dict[str, Any]],
    top_k: int = 5,
    text_field: str = "text",
) -> list[dict[str, Any]]:
    """
    Rerank candidates and return top_k.

    Default implementation using score(). Subclasses can override for optimization.
    """
    texts = [cand.get(text_field, cand.get("name", "")) for cand in candidates]
    scores = self.score(query, texts)

    scored = [
        {**candidate, "cross_encoder_score": float(score)}
        for candidate, score in zip(candidates, scores, strict=False)
    ]

    return heapq.nlargest(top_k, scored, key=lambda x: x["cross_encoder_score"])

novelentitymatcher.backends.sentencetransformer

Classes

Functions

novelentitymatcher.backends.litellm

Classes

novelentitymatcher.backends.static_embedding

Classes

StaticEmbeddingBackend(model_name, embedding_dim=None)

Bases: EmbeddingBackend

Backend for static embeddings.

Supports two approaches: 1. SentenceTransformer's StaticEmbedding module 2. model2vec StaticModel (for minishlab potion models and custom distillations)

Models: - RikkaBotan/stable-static-embedding-fast-retrieval-mrl-en (StaticEmbedding) - minishlab/potion-base-8M (model2vec) - minishlab/potion-base-32M (model2vec)

Parameters:

Name	Type	Description	Default
model_name	str	HuggingFace model name or local path	required
embedding_dim	int | None	Optional dimension for MRL models (None = full dimension)	None

Source code in src/novelentitymatcher/backends/static_embedding.py

def __init__(self, model_name: str, embedding_dim: int | None = None):
    """
    Initialize static embedding backend.

    Args:
        model_name: HuggingFace model name or local path
        embedding_dim: Optional dimension for MRL models (None = full dimension)
    """
    self.model_name = model_name
    self.embedding_dim = embedding_dim
    self.backend_type: str | None = None
    self.model: Any = None

    # Try to load with model2vec first (for minishlab models)
    try:
        from model2vec import StaticModel

        self.model = StaticModel.from_pretrained(model_name)
        self.backend_type = "model2vec"
        return
    except (ImportError, ValueError, RuntimeError):
        pass

    # Fall back to SentenceTransformer (for RikkaBotan MRL and others)
    try:
        # Some models (like RikkaBotan) require trust_remote_code=True
        # to load custom modules like SSE
        self.model = get_cached_sentence_transformer(
            model_name, trust_remote_code=True
        )
        self.backend_type = "sentence_transformers"
        self._is_native_static = any(
            isinstance(m, StaticEmbedding) for m in self.model.modules()
        )
        return
    except (OSError, RuntimeError, ValueError) as e:
        raise ValueError(
            f"Failed to load static embedding model {model_name}. "
            f"Tried both model2vec and SentenceTransformer. Last error: {e}"
        ) from e

Attributes

embedding_dimension property

Return the embedding dimension of the model.

Functions

encode(texts, batch_size=32)

Generate embeddings using static lookup.

Source code in src/novelentitymatcher/backends/static_embedding.py

def encode(self, texts, batch_size: int = 32):
    """Generate embeddings using static lookup."""
    if self.backend_type == "model2vec":
        embeddings = self.model.encode(texts)
    else:  # sentence_transformers
        embeddings = self.model.encode(
            texts,
            batch_size=batch_size,
            show_progress_bar=False,
        )

    # Convert to list if needed
    if hasattr(embeddings, "tolist"):
        embeddings = embeddings.tolist()
    elif not isinstance(embeddings, list):
        embeddings = list(embeddings)

    # Apply dimension reduction if requested (for MRL models)
    if self.embedding_dim is not None and embeddings:
        embeddings = [emb[: self.embedding_dim] for emb in embeddings]

    return embeddings

similarity(query_embeddings, corpus_embeddings, top_k)

Compute similarity using cosine similarity.

Source code in src/novelentitymatcher/backends/static_embedding.py

def similarity(self, query_embeddings, corpus_embeddings, top_k):
    """Compute similarity using cosine similarity."""
    query_embeddings = np.array(query_embeddings)
    corpus_embeddings = np.array(corpus_embeddings)

    similarities = cosine_similarity(query_embeddings, corpus_embeddings)

    results = []
    for _idx, query_sim in enumerate(similarities):
        top_indices = query_sim.argsort()[-top_k:][::-1]
        results.append(
            [
                {"corpus_id": int(i), "score": float(query_sim[i])}
                for i in top_indices
            ]
        )
    return results

Functions

novelentitymatcher.backends.reranker_st

SentenceTransformer-based cross-encoder reranker backend.

Classes

STReranker(model_name='BAAI/bge-reranker-v2-m3', device=None, batch_size=32)

Bases: RerankerBackend

SentenceTransformer cross-encoder reranker.

Uses CrossEncoder models for precise query-document scoring.

Parameters:

Name	Type	Description	Default
model_name	str	Name or path of the CrossEncoder model	'BAAI/bge-reranker-v2-m3'
device	str | None	Device to run model on (None for auto-detection)	None
batch_size	int	Batch size for inference	32

Source code in src/novelentitymatcher/backends/reranker_st.py

def __init__(
    self,
    model_name: str = "BAAI/bge-reranker-v2-m3",
    device: str | None = None,
    batch_size: int = 32,
):
    """
    Initialize the reranker.

    Args:
        model_name: Name or path of the CrossEncoder model
        device: Device to run model on (None for auto-detection)
        batch_size: Batch size for inference
    """
    self.model_name = model_name
    self.device = device
    self.batch_size = batch_size
    self.model = get_cached_cross_encoder(model_name, device=device)

Functions

score(query, docs)

Score query-document pairs.

Parameters:

Name	Type	Description	Default
query	str	Query text	required
docs	list[str]	List of document texts	required

Returns:

Type	Description
list[float]	List of scores (one per document)

Source code in src/novelentitymatcher/backends/reranker_st.py

def score(self, query: str, docs: list[str]) -> list[float]:
    """
    Score query-document pairs.

    Args:
        query: Query text
        docs: List of document texts

    Returns:
        List of scores (one per document)
    """
    # Prepare pairs
    pairs = [[query, doc] for doc in docs]

    # Score all pairs
    scores = self.model.predict(
        pairs,
        batch_size=self.batch_size,
        convert_to_numpy=False,
        convert_to_tensor=True,
    )

    # Convert to list of floats
    if hasattr(scores, "cpu"):
        scores = scores.cpu().tolist()
    elif hasattr(scores, "tolist"):
        scores = scores.tolist()
    elif not isinstance(scores, list):
        scores = list(scores)

    return scores

rerank(query, candidates, top_k=5, text_field='text')

Rerank candidates and return top_k.

Default implementation using score(). Subclasses can override for optimization.

Source code in src/novelentitymatcher/backends/base.py

def rerank(
    self,
    query: str,
    candidates: list[dict[str, Any]],
    top_k: int = 5,
    text_field: str = "text",
) -> list[dict[str, Any]]:
    """
    Rerank candidates and return top_k.

    Default implementation using score(). Subclasses can override for optimization.
    """
    texts = [cand.get(text_field, cand.get("name", "")) for cand in candidates]
    scores = self.score(query, texts)

    scored = [
        {**candidate, "cross_encoder_score": float(score)}
        for candidate, score in zip(candidates, scores, strict=False)
    ]

    return heapq.nlargest(top_k, scored, key=lambda x: x["cross_encoder_score"])

Functions