test-int/semantic/multilingual_benchmark.py
Rama main · réplica de lectura del repositorio original.
"""Typed benchmark cases, measurements, and artifacts for multilingual retrieval."""
from __future__ import annotations
import json
import math
import platform
import sys
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from statistics import mean
from typing import Literal, Sequence
import fastembed
import psutil
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncEngine
from basic_memory.config import BasicMemoryConfig, DatabaseBackend, default_fastembed_cache_dir
from basic_memory.repository.search_index_row import SearchIndexRow
from basic_memory.schemas.search import SearchRetrievalMode
from semantic.multilingual_corpus import MultilingualQuery, RetrievalCaseKind
@dataclass(frozen=True, slots=True)
class EmbeddingModelCase:
"""One FastEmbed model contract to benchmark in an isolated process."""
key: str
model_name: str
dimensions: int
cache_repository: str
model_file: str
catalog_size_gb: float
license: str
fallback_cache_directory: str | None = None
document_prefix: str | None = None
query_prefix: str | None = None
FASTEMBED_MODEL_CASES = (
EmbeddingModelCase(
key="bge-small-en",
model_name="BAAI/bge-small-en-v1.5",
dimensions=384,
cache_repository="qdrant/bge-small-en-v1.5-onnx-q",
model_file="model_optimized.onnx",
catalog_size_gb=0.067,
license="mit",
),
EmbeddingModelCase(
key="multilingual-minilm",
model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
dimensions=384,
cache_repository="qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
model_file="model_optimized.onnx",
catalog_size_gb=0.22,
license="apache-2.0",
),
EmbeddingModelCase(
key="multilingual-mpnet",
model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
dimensions=768,
cache_repository="xenova/paraphrase-multilingual-mpnet-base-v2",
model_file="onnx/model.onnx",
catalog_size_gb=1.0,
license="apache-2.0",
),
EmbeddingModelCase(
key="multilingual-e5-large",
model_name="intfloat/multilingual-e5-large",
dimensions=1024,
cache_repository="qdrant/multilingual-e5-large-onnx",
model_file="model.onnx",
catalog_size_gb=2.24,
license="mit",
fallback_cache_directory="fast-multilingual-e5-large",
document_prefix="passage: ",
query_prefix="query: ",
),
)
@dataclass(frozen=True, slots=True)
class BenchmarkStorageCase:
"""One database and vector-index pairing exercised by the benchmark."""
key: str
database_backend: DatabaseBackend
vector_index_name: Literal["sqlite-vec", "pgvector", "milvus"]
BENCHMARK_STORAGE_CASES = (
BenchmarkStorageCase("sqlite", DatabaseBackend.SQLITE, "sqlite-vec"),
BenchmarkStorageCase("postgres", DatabaseBackend.POSTGRES, "pgvector"),
BenchmarkStorageCase("milvus", DatabaseBackend.POSTGRES, "milvus"),
)
def embedding_model_case(key: str) -> EmbeddingModelCase:
"""Return one configured model case or fail with the supported keys."""
for model_case in FASTEMBED_MODEL_CASES:
if model_case.key == key:
return model_case
supported = ", ".join(model.key for model in FASTEMBED_MODEL_CASES)
raise ValueError(f"Unknown multilingual benchmark model {key!r}; choose one of: {supported}")
def benchmark_storage_case(value: str) -> BenchmarkStorageCase:
"""Parse a supported SQL database and vector-index pairing."""
normalized = value.strip().lower()
for storage_case in BENCHMARK_STORAGE_CASES:
if storage_case.key == normalized:
return storage_case
raise ValueError("Multilingual benchmark backend must be 'sqlite', 'postgres', or 'milvus'")
def benchmark_retrieval_mode(value: str) -> SearchRetrievalMode:
"""Parse the retrieval mode supported by the embedding benchmark."""
try:
mode = SearchRetrievalMode(value.strip().lower())
except ValueError as exc:
raise ValueError("Multilingual benchmark mode must be 'vector' or 'hybrid'") from exc
if mode not in {SearchRetrievalMode.VECTOR, SearchRetrievalMode.HYBRID}:
raise ValueError("Multilingual benchmark mode must be 'vector' or 'hybrid'")
return mode
def embedding_benchmark_config(
model_case: EmbeddingModelCase,
storage_case: BenchmarkStorageCase,
*,
milvus_uri: str | None = None,
) -> BasicMemoryConfig:
"""Build the Cloud-like FastEmbed configuration used by benchmark services."""
if storage_case.vector_index_name == "milvus" and milvus_uri is None:
raise ValueError("The Milvus benchmark requires a local or remote Milvus URI")
configured_vector_index: Literal["pgvector", "milvus"] = (
"milvus" if storage_case.vector_index_name == "milvus" else "pgvector"
)
return BasicMemoryConfig(
env="test",
database_backend=storage_case.database_backend,
semantic_vector_index=configured_vector_index,
semantic_embedding_provider="fastembed",
semantic_embedding_model=model_case.model_name,
semantic_embedding_dimensions=model_case.dimensions,
semantic_embedding_document_prefix=model_case.document_prefix,
semantic_embedding_query_prefix=model_case.query_prefix,
semantic_embedding_batch_size=8,
semantic_embedding_sync_batch_size=64,
semantic_embedding_threads=1,
semantic_embedding_parallel=1,
semantic_min_similarity=0.55,
milvus_uri=milvus_uri,
)
@dataclass(frozen=True, slots=True)
class RetrievalObservation:
"""The ranked and production-threshold outcomes for one query."""
query: MultilingualQuery
ranked_permalinks: tuple[str, ...]
accepted_permalinks: tuple[str, ...]
ranking_latency_seconds: float
accepted_latency_seconds: float
@property
def first_relevant_rank(self) -> int | None:
relevant = set(self.query.relevant_permalinks)
for rank, permalink in enumerate(self.ranked_permalinks[:10], start=1):
if permalink in relevant:
return rank
return None
@dataclass(frozen=True, slots=True)
class RetrievalSummary:
"""Aggregate retrieval quality for a named corpus slice."""
name: str
query_count: int
positive_query_count: int
negative_query_count: int
recall_at_5: float
mrr_at_10: float
accepted_empty_rate: float
wrong_top_rate: float
negative_false_positive_rate: float
ranking_p50_ms: float
ranking_p95_ms: float
accepted_p50_ms: float
accepted_p95_ms: float
def as_metrics(self) -> dict[str, float | int]:
"""Return numeric metrics compatible with the existing JSONL comparator."""
return {
"queries_executed": self.query_count,
"positive_queries": self.positive_query_count,
"negative_queries": self.negative_query_count,
"recall_at_5": round(self.recall_at_5, 4),
"mrr_at_10": round(self.mrr_at_10, 4),
"accepted_empty_rate": round(self.accepted_empty_rate, 4),
"wrong_top_rate": round(self.wrong_top_rate, 4),
"negative_false_positive_rate": round(self.negative_false_positive_rate, 4),
"ranking_p50_ms": round(self.ranking_p50_ms, 2),
"ranking_p95_ms": round(self.ranking_p95_ms, 2),
"accepted_p50_ms": round(self.accepted_p50_ms, 2),
"accepted_p95_ms": round(self.accepted_p95_ms, 2),
}
@dataclass(frozen=True, slots=True)
class RuntimeMeasurements:
"""Operational measurements captured by one isolated model/backend run."""
document_count: int
cold_load_seconds: float
indexing_seconds: float
rss_before_load_bytes: int
rss_after_load_bytes: int
rss_after_index_bytes: int
peak_rss_bytes: int
model_cache_bytes: int
vector_storage_bytes: int
def as_metrics(self) -> dict[str, float | int]:
"""Return comparable numeric runtime metrics."""
documents_per_second = (
self.document_count / self.indexing_seconds if self.indexing_seconds else 0.0
)
return {
"documents_indexed": self.document_count,
"cold_load_seconds": round(self.cold_load_seconds, 4),
"indexing_seconds": round(self.indexing_seconds, 4),
"documents_per_sec": round(documents_per_second, 4),
"rss_before_load_bytes": self.rss_before_load_bytes,
"rss_after_load_bytes": self.rss_after_load_bytes,
"rss_model_delta_bytes": max(0, self.rss_after_load_bytes - self.rss_before_load_bytes),
"rss_after_index_bytes": self.rss_after_index_bytes,
"peak_rss_bytes": self.peak_rss_bytes,
"model_cache_bytes": self.model_cache_bytes,
"vector_storage_bytes": self.vector_storage_bytes,
}
def result_permalinks(
results: Sequence[SearchIndexRow],
query: MultilingualQuery,
) -> tuple[str, ...]:
"""Return ranked permalinks, invalidating only a target that fails its chunk oracle."""
relevant_permalinks = set(query.relevant_permalinks)
permalinks: list[str] = []
for result in results:
if result.permalink is None:
continue
if query.required_chunk_text is not None and result.permalink in relevant_permalinks:
top_chunk = (result.matched_chunk_text or "").partition("\n---\n")[0]
if query.required_chunk_text not in top_chunk:
continue
permalinks.append(result.permalink)
return tuple(permalinks)
def percentile_ms(samples: Sequence[float], quantile: float) -> float:
"""Return a linearly interpolated percentile in milliseconds."""
if not samples:
return 0.0
ordered = sorted(samples)
position = (len(ordered) - 1) * quantile
lower_index = math.floor(position)
upper_index = math.ceil(position)
if lower_index == upper_index:
return ordered[lower_index] * 1000
fraction = position - lower_index
value = ordered[lower_index] + (ordered[upper_index] - ordered[lower_index]) * fraction
return value * 1000
def summarize_retrieval(
name: str,
observations: Sequence[RetrievalObservation],
) -> RetrievalSummary:
"""Compute quality, threshold, and latency metrics for one corpus slice."""
positive = [item for item in observations if item.query.relevant_permalinks]
negative = [item for item in observations if not item.query.relevant_permalinks]
recalls: list[float] = []
reciprocal_ranks: list[float] = []
for observation in positive:
relevant = set(observation.query.relevant_permalinks)
retrieved = set(observation.ranked_permalinks[:5])
recalls.append(len(relevant.intersection(retrieved)) / len(relevant))
rank = observation.first_relevant_rank
reciprocal_ranks.append(1.0 / rank if rank is not None else 0.0)
accepted_empty_rate = (
sum(not item.accepted_permalinks for item in positive) / len(positive) if positive else 0.0
)
wrong_top_rate = (
sum(
not item.ranked_permalinks
or item.ranked_permalinks[0] not in item.query.relevant_permalinks
for item in positive
)
/ len(positive)
if positive
else 0.0
)
false_positive_rate = (
sum(bool(item.accepted_permalinks) for item in negative) / len(negative)
if negative
else 0.0
)
ranking_latencies = [item.ranking_latency_seconds for item in observations]
accepted_latencies = [item.accepted_latency_seconds for item in observations]
return RetrievalSummary(
name=name,
query_count=len(observations),
positive_query_count=len(positive),
negative_query_count=len(negative),
recall_at_5=mean(recalls) if recalls else 0.0,
mrr_at_10=mean(reciprocal_ranks) if reciprocal_ranks else 0.0,
accepted_empty_rate=accepted_empty_rate,
wrong_top_rate=wrong_top_rate,
negative_false_positive_rate=false_positive_rate,
ranking_p50_ms=percentile_ms(ranking_latencies, 0.50),
ranking_p95_ms=percentile_ms(ranking_latencies, 0.95),
accepted_p50_ms=percentile_ms(accepted_latencies, 0.50),
accepted_p95_ms=percentile_ms(accepted_latencies, 0.95),
)
def retrieval_summaries(
observations: Sequence[RetrievalObservation],
) -> tuple[RetrievalSummary, ...]:
"""Build overall, behavior, and language slices from one benchmark run."""
summaries = [summarize_retrieval("overall", observations)]
for kind in RetrievalCaseKind:
kind_observations = [item for item in observations if item.query.kind is kind]
if kind_observations:
summaries.append(summarize_retrieval(kind.value, kind_observations))
for language in sorted({item.query.language for item in observations}):
language_observations = [item for item in observations if item.query.language == language]
summaries.append(summarize_retrieval(f"language-{language}", language_observations))
return tuple(summaries)
def process_rss_bytes() -> int:
"""Return the process's current resident memory."""
return psutil.Process().memory_info().rss
def process_peak_rss_bytes(*, platform_name: str = sys.platform) -> int:
"""Return the process high-water resident memory across supported platforms."""
if platform_name == "win32":
return int(psutil.Process().memory_info().peak_wset)
import resource
peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
return int(peak if platform_name == "darwin" else peak * 1024)
def model_cache_size_bytes(
model_case: EmbeddingModelCase,
benchmark_config: BasicMemoryConfig,
) -> int:
"""Measure the selected model's materialized FastEmbed cache subtree."""
configured_cache = benchmark_config.semantic_embedding_cache_dir
cache_root = Path(configured_cache or default_fastembed_cache_dir())
model_directory = cache_root / f"models--{model_case.cache_repository.replace('/', '--')}"
snapshots = model_directory / "snapshots"
if any(snapshots.glob(f"*/{model_case.model_file}")):
return directory_size_bytes(model_directory)
if model_case.fallback_cache_directory is None:
return 0
return directory_size_bytes(cache_root / model_case.fallback_cache_directory)
def directory_size_bytes(directory: Path) -> int:
"""Measure unique file bytes without double-counting Hugging Face hardlinks."""
if not directory.exists():
return 0
seen_files: set[tuple[int, int]] = set()
total = 0
for path in directory.rglob("*"):
if not path.is_file():
continue
stat = path.stat()
identity = (stat.st_dev, stat.st_ino)
if identity in seen_files:
continue
seen_files.add(identity)
total += stat.st_size
return total
async def vector_storage_size_bytes(
engine: AsyncEngine,
storage_case: BenchmarkStorageCase,
*,
milvus_storage_directory: Path | None = None,
) -> int:
"""Measure database storage after the corpus vectors have been indexed."""
external_vector_bytes = 0
if storage_case.vector_index_name == "milvus":
if milvus_storage_directory is None:
raise ValueError("Milvus storage measurement requires its local storage directory")
external_vector_bytes = directory_size_bytes(milvus_storage_directory)
if storage_case.database_backend is DatabaseBackend.SQLITE:
async with engine.connect() as connection:
result = await connection.execute(
text(
"SELECT COALESCE(SUM(pgsize), 0) FROM dbstat "
"WHERE name = 'search_vector_chunks' "
"OR name IN ("
"SELECT name FROM sqlite_schema WHERE tbl_name = 'search_vector_chunks'"
") "
"OR name LIKE 'search_vector_embeddings%' "
"OR name IN ("
"SELECT name FROM sqlite_schema "
"WHERE tbl_name LIKE 'search_vector_embeddings%'"
")"
)
)
return int(result.scalar_one())
async with engine.connect() as connection:
relations = (
"pg_total_relation_size('search_vector_chunks')"
if storage_case.vector_index_name == "milvus"
else "pg_total_relation_size('search_vector_chunks') + "
"pg_total_relation_size('search_vector_embeddings')"
)
result = await connection.execute(text(f"SELECT {relations}"))
return external_vector_bytes + int(result.scalar_one())
def benchmark_context(
model_case: EmbeddingModelCase,
storage_case: BenchmarkStorageCase,
retrieval_mode: SearchRetrievalMode,
corpus_version: str,
threshold: float,
) -> dict[str, object]:
"""Describe the software, model, and host context for an artifact."""
return {
"model_key": model_case.key,
"model_name": model_case.model_name,
"dimensions": model_case.dimensions,
"document_prefix": model_case.document_prefix,
"query_prefix": model_case.query_prefix,
"license": model_case.license,
"catalog_size_gb": model_case.catalog_size_gb,
"backend": storage_case.key,
"database_backend": storage_case.database_backend.value,
"vector_index": storage_case.vector_index_name,
"retrieval_mode": retrieval_mode.value,
"corpus_version": corpus_version,
"production_similarity_threshold": threshold,
"python_version": platform.python_version(),
"fastembed_version": fastembed.__version__,
"platform": platform.platform(),
"processor": platform.processor(),
"logical_cpu_count": psutil.cpu_count(),
"physical_memory_bytes": psutil.virtual_memory().total,
}
def write_multilingual_benchmark_artifact(
output_path: Path,
context: dict[str, object],
summaries: Sequence[RetrievalSummary],
runtime: RuntimeMeasurements,
observations: Sequence[RetrievalObservation],
) -> None:
"""Write comparison-compatible JSONL records for one isolated run."""
output_path.parent.mkdir(parents=True, exist_ok=True)
timestamp = datetime.now(timezone.utc).isoformat()
observation_payload = [
{
"query": item.query.name,
"kind": item.query.kind.value,
"language": item.query.language,
"relevant": list(item.query.relevant_permalinks),
"ranked": list(item.ranked_permalinks),
"accepted": list(item.accepted_permalinks),
"first_relevant_rank": item.first_relevant_rank,
}
for item in observations
]
with output_path.open("a", encoding="utf-8") as artifact:
for summary in summaries:
payload: dict[str, object] = {
"benchmark": f"multilingual-quality-{summary.name}",
"timestamp_utc": timestamp,
"context": context,
"metrics": summary.as_metrics(),
}
if summary.name == "overall":
payload["observations"] = observation_payload
artifact.write(json.dumps(payload, ensure_ascii=False, sort_keys=True) + "\n")
artifact.write(
json.dumps(
{
"benchmark": "multilingual-runtime",
"timestamp_utc": timestamp,
"context": context,
"metrics": runtime.as_metrics(),
},
ensure_ascii=False,
sort_keys=True,
)
+ "\n"
)