embedflow 0.1.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- embedflow/__init__.py +25 -0
- embedflow/__main__.py +3 -0
- embedflow/analysis.py +192 -0
- embedflow/cache/__init__.py +4 -0
- embedflow/cache/base.py +28 -0
- embedflow/cache/persistent_cache.py +198 -0
- embedflow/cli.py +1200 -0
- embedflow/compatibility/__init__.py +28 -0
- embedflow/compatibility/candidate_gap.py +105 -0
- embedflow/compatibility/containment.py +17 -0
- embedflow/compatibility/evaluate.py +319 -0
- embedflow/compatibility/metrics.py +75 -0
- embedflow/compatibility/migration_depth.py +67 -0
- embedflow/compatibility/probe.py +34 -0
- embedflow/compatibility/report.py +102 -0
- embedflow/compatibility/t2.py +64 -0
- embedflow/config.py +455 -0
- embedflow/data/__init__.py +1 -0
- embedflow/data/registry/__init__.py +1 -0
- embedflow/data/registry/benchmark_profiles.jsonl +3 -0
- embedflow/data/registry/checksums.sha256 +4 -0
- embedflow/data/registry/migrations.jsonl +15 -0
- embedflow/data/registry/registry_manifest.json +16 -0
- embedflow/data/registry/research_summaries.json +55 -0
- embedflow/data/registry/schema_version.json +5 -0
- embedflow/frozen/T2_V1_FROZEN_SPEC.md +71 -0
- embedflow/frozen/T2_V1_FROZEN_SPEC.sha256 +1 -0
- embedflow/indexes/__init__.py +5 -0
- embedflow/indexes/base.py +60 -0
- embedflow/indexes/faiss_backend.py +240 -0
- embedflow/indexes/qdrant_backend.py +225 -0
- embedflow/metrics/__init__.py +3 -0
- embedflow/metrics/latency.py +50 -0
- embedflow/migration/__init__.py +3 -0
- embedflow/migration/compatibility.py +156 -0
- embedflow/migration/facade.py +312 -0
- embedflow/migration/materializer.py +190 -0
- embedflow/migration/planner.py +78 -0
- embedflow/migration/state.py +81 -0
- embedflow/models/__init__.py +4 -0
- embedflow/models/base.py +31 -0
- embedflow/models/huggingface.py +226 -0
- embedflow/registry/__init__.py +47 -0
- embedflow/registry/loader.py +785 -0
- embedflow/registry/matcher.py +197 -0
- embedflow/registry/schema.py +266 -0
- embedflow/runtime.py +115 -0
- embedflow/serving/__init__.py +3 -0
- embedflow/serving/api.py +161 -0
- embedflow/serving/engine.py +222 -0
- embedflow/serving/factory.py +3 -0
- embedflow/serving/schemas.py +39 -0
- embedflow-0.1.0.dist-info/METADATA +210 -0
- embedflow-0.1.0.dist-info/RECORD +64 -0
- embedflow-0.1.0.dist-info/WHEEL +5 -0
- embedflow-0.1.0.dist-info/entry_points.txt +2 -0
- embedflow-0.1.0.dist-info/licenses/LICENSE +178 -0
- embedflow-0.1.0.dist-info/top_level.txt +2 -0
- src/__init__.py +1 -0
- src/embed.py +123 -0
- src/probe_features.py +24 -0
- src/storage.py +51 -0
- src/t2_v1.py +21 -0
- src/utils.py +53 -0
embedflow/serving/api.py
ADDED
|
@@ -0,0 +1,161 @@
|
|
|
1
|
+
import html
|
|
2
|
+
import math
|
|
3
|
+
from typing import Any
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
def dashboard_html(title: str = "EmbedFlow") -> str:
|
|
7
|
+
# Keep this as an ordinary template instead of a Python f-string. The
|
|
8
|
+
# dashboard contains JavaScript template literals, and escaping every
|
|
9
|
+
# JavaScript brace in a f-string made browser failures very easy to hide.
|
|
10
|
+
escaped_title = html.escape(title, quote=True)
|
|
11
|
+
return '''<!doctype html><html><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1"><title>__TITLE__</title>
|
|
12
|
+
<style>
|
|
13
|
+
body{font:15px system-ui,sans-serif;background:#0b1020;color:#e8eefc;max-width:1100px;margin:0 auto;padding:28px}
|
|
14
|
+
.card{background:#141d35;border:1px solid #2d3b62;border-radius:14px;padding:18px;margin:12px 0}
|
|
15
|
+
.grid{display:grid;grid-template-columns:repeat(auto-fit,minmax(210px,1fr));gap:12px}
|
|
16
|
+
.label{color:#9fb0d8;font-size:12px;text-transform:uppercase;letter-spacing:.08em}
|
|
17
|
+
.value{font-size:25px;font-weight:650;margin-top:5px}
|
|
18
|
+
input{width:min(75%,650px);padding:12px;border-radius:8px;border:1px solid #45567f;background:#0b1020;color:white}
|
|
19
|
+
button{padding:12px 18px;border:0;border-radius:8px;background:#6ea8fe;color:#071021;font-weight:700;cursor:pointer}
|
|
20
|
+
button:disabled{opacity:.6;cursor:wait}
|
|
21
|
+
table{width:100%;border-collapse:collapse;margin-top:12px}td,th{padding:8px;border-bottom:1px solid #2d3b62;text-align:left}code{color:#a8d5ff}
|
|
22
|
+
.error{color:#ffb4b4;background:#3a1820;border-color:#7c3343}.muted{color:#9fb0d8}
|
|
23
|
+
</style></head>
|
|
24
|
+
<body><h1>EMBEDFLOW</h1><p>Progressive Embedding Migration</p><div id="status" class="card">Connecting to EmbedFlow…</div>
|
|
25
|
+
<div class="card"><h2>Search</h2><form id="form"><input id="query" placeholder="Ask a question…" autocomplete="off" required><button id="submit" type="submit">Search</button></form><div id="search" class="muted">Enter a query to test source retrieval, target reranking, and progressive caching.</div></div>
|
|
26
|
+
<script>
|
|
27
|
+
const fmt=(x)=>x===undefined||x===null?'—':(typeof x==='number'?x.toFixed(2):x);
|
|
28
|
+
const esc=(x)=>String(x??'').replace(/[&<>"']/g,(c)=>({'&':'&','<':'<','>':'>','"':'"',"'":'''}[c]));
|
|
29
|
+
async function jsonFetch(url, options){
|
|
30
|
+
const response=await fetch(url, options);
|
|
31
|
+
const text=await response.text();
|
|
32
|
+
let data={};
|
|
33
|
+
try{data=text?JSON.parse(text):{};}catch(_){data={detail:text||response.statusText};}
|
|
34
|
+
if(!response.ok){
|
|
35
|
+
const detail=data&&data.detail;
|
|
36
|
+
const message=Array.isArray(detail)?detail.map((item)=>typeof item==='string'?item:(item.msg||JSON.stringify(item))).join('; '):(detail&&typeof detail==='object'?JSON.stringify(detail):(detail||('HTTP '+response.status)));
|
|
37
|
+
throw new Error(message);
|
|
38
|
+
}
|
|
39
|
+
return data;
|
|
40
|
+
}
|
|
41
|
+
async function refresh(){
|
|
42
|
+
try{
|
|
43
|
+
const s=await jsonFetch('/status');
|
|
44
|
+
const m=s.migration||{},c=s.cache||{},w=s.worker||{},l=s.latency||{};
|
|
45
|
+
document.querySelector('#status').className='card';
|
|
46
|
+
document.querySelector('#status').innerHTML=`<div class="grid"><div><div class="label">Source</div><div class="value"><code>${esc(m.source_model||'configured')}</code></div></div><div><div class="label">Target</div><div class="value"><code>${esc(m.target_model||'configured')}</code></div></div><div><div class="label">Diagnostic</div><div class="value">${esc(m.diagnostic)}</div></div><div><div class="label">Candidate depth</div><div class="value">K=${fmt(m.candidate_depth)}</div></div><div><div class="label">Cache progress</div><div class="value">${fmt(c.cached_target_vectors)} / ${fmt(m.corpus_size)} (${fmt(100*(m.cache_fraction||0))}%)</div></div><div><div class="label">Background throughput</div><div class="value">${fmt(w.last_throughput_docs_sec)} docs/s</div></div><div><div class="label">Query p50/p95</div><div class="value">${fmt(l.p50_ms)} / ${fmt(l.p95_ms)} ms</div></div></div><p>Status: <b>${esc(m.status)}</b> · ANN: <b>${esc(m.ann_status)}</b></p>`;
|
|
47
|
+
}catch(error){
|
|
48
|
+
document.querySelector('#status').className='card error';
|
|
49
|
+
document.querySelector('#status').textContent='Unable to load migration status: '+error.message;
|
|
50
|
+
}
|
|
51
|
+
}
|
|
52
|
+
refresh();
|
|
53
|
+
setInterval(refresh,3000);
|
|
54
|
+
document.querySelector('#form').addEventListener('submit',async(event)=>{
|
|
55
|
+
event.preventDefault();
|
|
56
|
+
const input=document.querySelector('#query'),button=document.querySelector('#submit'),out=document.querySelector('#search');
|
|
57
|
+
const query=input.value.trim(); if(!query)return;
|
|
58
|
+
button.disabled=true; button.textContent='Searching…'; out.className='muted'; out.textContent='Running source retrieval and target reranking…';
|
|
59
|
+
try{
|
|
60
|
+
const x=await jsonFetch('/search',{method:'POST',headers:{'content-type':'application/json'},body:JSON.stringify({query:query,top_k:10})});
|
|
61
|
+
const migration=x.migration||{}, timing=x.timing_ms||{};
|
|
62
|
+
out.className='';
|
|
63
|
+
out.innerHTML=`<p><b>${esc(migration.status)}</b> · ${fmt(timing.total_ms)} ms · cache ${fmt(migration.target_cache_hits)}/${fmt(migration.source_candidates)} · queued ${fmt(migration.async_misses_queued)}</p><table><tr><th>Target rank</th><th>Source rank</th><th>Cached</th><th>Document</th></tr>${(x.results||[]).map((a)=>`<tr><td>${fmt(a.target_rank)}</td><td>${fmt(a.source_rank)}</td><td>${a.target_vector_cached?'yes':'no'}</td><td>${esc(a.text)}</td></tr>`).join('')}</table><p class="muted">Total ${fmt(timing.total_ms)} ms · source encode ${fmt(timing.source_query_encode_ms)} ms · ANN ${fmt(timing.source_ann_ms)} ms · target encode ${fmt(timing.target_query_encode_ms)} ms · scoring ${fmt(timing.target_score_ms)} ms</p>`;
|
|
64
|
+
}catch(error){
|
|
65
|
+
out.className='card error'; out.textContent='Search failed: '+error.message+' (check the server terminal for details)';
|
|
66
|
+
}finally{button.disabled=false;button.textContent='Search';}
|
|
67
|
+
});
|
|
68
|
+
</script></body></html>'''.replace('__TITLE__', escaped_title)
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
def create_app(engine: Any):
|
|
72
|
+
try:
|
|
73
|
+
from fastapi import FastAPI, HTTPException
|
|
74
|
+
from fastapi.responses import HTMLResponse
|
|
75
|
+
|
|
76
|
+
from .schemas import PrewarmRequest, SearchRequest, SearchResponse
|
|
77
|
+
except ImportError as exc:
|
|
78
|
+
raise RuntimeError("API requires fastapi, uvicorn, and pydantic") from exc
|
|
79
|
+
|
|
80
|
+
app = FastAPI(title="EmbedFlow", version="0.1.0")
|
|
81
|
+
|
|
82
|
+
@app.get("/", response_class=HTMLResponse)
|
|
83
|
+
def root(): return dashboard_html(engine.cfg.dashboard_title)
|
|
84
|
+
|
|
85
|
+
@app.get("/health")
|
|
86
|
+
def health(): return {"status": "ok"}
|
|
87
|
+
|
|
88
|
+
@app.get("/status")
|
|
89
|
+
def status(): return engine.status()
|
|
90
|
+
|
|
91
|
+
@app.get("/plan")
|
|
92
|
+
def plan(): return engine.plan.to_dict()
|
|
93
|
+
|
|
94
|
+
@app.post("/analyze")
|
|
95
|
+
def analyze():
|
|
96
|
+
"""Return the persisted finite-tail decision and deployment plan."""
|
|
97
|
+
return {
|
|
98
|
+
"diagnostic": engine.plan.diagnostic,
|
|
99
|
+
"recommended_initial_k": engine.plan.candidate_depth,
|
|
100
|
+
"ann_status": engine.plan.ann_status,
|
|
101
|
+
"plan": engine.plan.to_dict(),
|
|
102
|
+
"probe": engine.probe,
|
|
103
|
+
"warning": "T2-v1 is an empirical finite-tail diagnostic, not a compatibility guarantee.",
|
|
104
|
+
}
|
|
105
|
+
|
|
106
|
+
@app.post("/search", response_model=SearchResponse)
|
|
107
|
+
def search(body: SearchRequest):
|
|
108
|
+
try:
|
|
109
|
+
result = engine.search(body.query, body.top_k, body.candidate_depth, body.max_sync_misses)
|
|
110
|
+
migration = result.get("migration", {})
|
|
111
|
+
# Keep the detailed nested object and expose the compact fields
|
|
112
|
+
# shown in the public API example at the top level as well.
|
|
113
|
+
result.update({
|
|
114
|
+
"state": migration.get("status"),
|
|
115
|
+
"candidate_depth": migration.get("candidate_depth"),
|
|
116
|
+
"cache_hits": migration.get("cache_hits", migration.get("target_cache_hits")),
|
|
117
|
+
"cache_misses": migration.get("cache_misses", migration.get("target_cache_misses")),
|
|
118
|
+
"sync_encoded": migration.get("sync_encoded", migration.get("sync_misses_encoded")),
|
|
119
|
+
"async_queued": migration.get("async_queued", migration.get("async_misses_queued")),
|
|
120
|
+
})
|
|
121
|
+
return result
|
|
122
|
+
except (ValueError, KeyError, RuntimeError, TypeError, OSError) as exc:
|
|
123
|
+
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
124
|
+
except Exception as exc: # keep backend/model failures as structured API errors
|
|
125
|
+
raise HTTPException(status_code=500, detail=f"search backend failure: {exc}") from exc
|
|
126
|
+
|
|
127
|
+
@app.post("/prewarm")
|
|
128
|
+
def prewarm(body: PrewarmRequest):
|
|
129
|
+
try:
|
|
130
|
+
return engine.prewarm(body.document_ids, body.asynchronous)
|
|
131
|
+
except (ValueError, KeyError, RuntimeError, TypeError, OSError) as exc:
|
|
132
|
+
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
133
|
+
except Exception as exc:
|
|
134
|
+
raise HTTPException(status_code=500, detail=f"prewarm backend failure: {exc}") from exc
|
|
135
|
+
|
|
136
|
+
@app.get("/metrics")
|
|
137
|
+
def metrics():
|
|
138
|
+
from ..metrics import aggregate_records, summarize
|
|
139
|
+
|
|
140
|
+
values: list[float] = []
|
|
141
|
+
for row in engine.records:
|
|
142
|
+
try:
|
|
143
|
+
value = float(row.get("total_ms"))
|
|
144
|
+
except (TypeError, ValueError):
|
|
145
|
+
continue
|
|
146
|
+
if math.isfinite(value) and value >= 0:
|
|
147
|
+
values.append(value)
|
|
148
|
+
return {
|
|
149
|
+
"latency": summarize(values) if values else {"count": 0},
|
|
150
|
+
"stages": aggregate_records(engine.records),
|
|
151
|
+
"cache": engine.cache.stats(),
|
|
152
|
+
"queue": engine.worker.stats(),
|
|
153
|
+
}
|
|
154
|
+
|
|
155
|
+
@app.get("/economics")
|
|
156
|
+
def economics():
|
|
157
|
+
from ..cli import economics_for
|
|
158
|
+
return economics_for(engine.documents.size(), engine.cfg.economics.target_docs_per_second, engine.cfg.economics.gpu_price_per_hour,
|
|
159
|
+
engine.cache.stats()["cached_target_vectors"])
|
|
160
|
+
|
|
161
|
+
return app
|
|
@@ -0,0 +1,222 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
import threading
|
|
5
|
+
import time
|
|
6
|
+
from pathlib import Path
|
|
7
|
+
from typing import Any
|
|
8
|
+
|
|
9
|
+
import numpy as np
|
|
10
|
+
|
|
11
|
+
from ..cache import SQLiteVectorCache
|
|
12
|
+
from ..metrics import summarize
|
|
13
|
+
from ..migration.materializer import MaterializationWorker
|
|
14
|
+
from ..migration.planner import MigrationPlan, make_plan
|
|
15
|
+
from ..migration.state import DocumentStore, MigrationState
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def _elapsed(start: int) -> float:
|
|
19
|
+
return (time.perf_counter_ns() - start) / 1_000_000.0
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class MigrationEngine:
|
|
23
|
+
"""Backend-agnostic source retrieval + target reranking serving engine."""
|
|
24
|
+
|
|
25
|
+
def __init__(self, cfg: Any, source_model: Any, target_model: Any, source_index: Any,
|
|
26
|
+
target_cache: SQLiteVectorCache, documents: DocumentStore,
|
|
27
|
+
target_index: Any | None = None, probe: dict[str, Any] | None = None,
|
|
28
|
+
ann_status: str = "UNKNOWN", start_worker: bool = True):
|
|
29
|
+
self.cfg, self.source_model, self.target_model = cfg, source_model, target_model
|
|
30
|
+
self.source_index, self.target_index, self.cache, self.documents = source_index, target_index, target_cache, documents
|
|
31
|
+
self.state = MigrationState(cfg.state_path, cfg)
|
|
32
|
+
self.probe = probe or {}
|
|
33
|
+
self.ann_status = ann_status
|
|
34
|
+
self.plan: MigrationPlan = make_plan(target_model=cfg.target.model, corpus_size=documents.size(),
|
|
35
|
+
diagnostic=self.probe.get("diagnostic", "UNKNOWN"), recommended_k=self.probe.get("recommended_k"),
|
|
36
|
+
ann_status=ann_status, cached_target_vectors=target_cache.stats()["cached_target_vectors"],
|
|
37
|
+
default_k=cfg.migration.candidate_depth, throughput_docs_per_second=cfg.economics.target_docs_per_second,
|
|
38
|
+
gpu_price_per_hour=cfg.economics.gpu_price_per_hour, source_model=cfg.source.model)
|
|
39
|
+
queue_path = Path(cfg.cache.path) / "materialization_queue.sqlite3"
|
|
40
|
+
self.worker = MaterializationWorker(target_model, documents, target_cache, queue_path,
|
|
41
|
+
batch_size=cfg.migration.background_batch_size,
|
|
42
|
+
max_retries=cfg.migration.max_retries, state=self.state)
|
|
43
|
+
self.records: list[dict[str, Any]] = []
|
|
44
|
+
telemetry_path = getattr(getattr(cfg, "telemetry", None), "latency_log", None)
|
|
45
|
+
self.metrics_path = Path(telemetry_path) if telemetry_path else Path(cfg.cache.path) / "latency.jsonl"
|
|
46
|
+
if self.metrics_path.exists():
|
|
47
|
+
try:
|
|
48
|
+
with self.metrics_path.open() as handle:
|
|
49
|
+
loaded = [json.loads(line) for line in handle if line.strip()]
|
|
50
|
+
self.records = [row for row in loaded if isinstance(row, dict)]
|
|
51
|
+
except (OSError, json.JSONDecodeError):
|
|
52
|
+
# A malformed metrics tail must not make cached vectors
|
|
53
|
+
# unusable; fresh measurements continue in a new process.
|
|
54
|
+
self.records = []
|
|
55
|
+
self._lock = threading.RLock()
|
|
56
|
+
self._closed = False
|
|
57
|
+
self.state.update(status="SERVING", diagnostic=self.plan.diagnostic, ann_status=ann_status,
|
|
58
|
+
candidate_depth=self.plan.candidate_depth, corpus_size=documents.size(),
|
|
59
|
+
migration_strategy=self.plan.migration_strategy)
|
|
60
|
+
if start_worker: self.worker.start()
|
|
61
|
+
|
|
62
|
+
def _ensure_open(self) -> None:
|
|
63
|
+
if self._closed:
|
|
64
|
+
raise RuntimeError("migration engine is closed")
|
|
65
|
+
|
|
66
|
+
def _record(self, row: dict[str, Any]) -> None:
|
|
67
|
+
with self._lock:
|
|
68
|
+
self.records.append(row)
|
|
69
|
+
self.metrics_path.parent.mkdir(parents=True, exist_ok=True)
|
|
70
|
+
with self.metrics_path.open("a") as handle:
|
|
71
|
+
handle.write(json.dumps(row, ensure_ascii=False, default=float) + "\n")
|
|
72
|
+
handle.flush()
|
|
73
|
+
self.state.update(queries=int(self.state.snapshot().get("queries", 0)) + 1,
|
|
74
|
+
last_latency_ms=row["total_ms"], cache=self.cache.stats(), worker=self.worker.stats())
|
|
75
|
+
|
|
76
|
+
def search(self, query: str, top_k: int = 10, candidate_depth: int | None = None,
|
|
77
|
+
max_sync_misses: int | None = None) -> dict[str, Any]:
|
|
78
|
+
self._ensure_open()
|
|
79
|
+
if not isinstance(query, str) or not query.strip(): raise ValueError("query must be a non-empty string")
|
|
80
|
+
if isinstance(top_k, bool) or int(top_k) != top_k or int(top_k) < 1:
|
|
81
|
+
raise ValueError("top_k must be a positive integer")
|
|
82
|
+
if candidate_depth is not None and (isinstance(candidate_depth, bool) or int(candidate_depth) != candidate_depth or int(candidate_depth) < 1):
|
|
83
|
+
raise ValueError("candidate_depth must be a positive integer")
|
|
84
|
+
if max_sync_misses is not None and (isinstance(max_sync_misses, bool) or int(max_sync_misses) != max_sync_misses or int(max_sync_misses) < 0):
|
|
85
|
+
raise ValueError("max_sync_misses must be a non-negative integer")
|
|
86
|
+
K = int(candidate_depth or self.plan.candidate_depth or self.cfg.migration.candidate_depth)
|
|
87
|
+
if K < 1:
|
|
88
|
+
raise ValueError("candidate_depth must be a positive integer")
|
|
89
|
+
max_sync = self.cfg.migration.max_sync_misses if max_sync_misses is None else int(max_sync_misses)
|
|
90
|
+
total_start = time.perf_counter_ns(); row: dict[str, Any] = {"mode": "embedflow", "K": K}
|
|
91
|
+
t = time.perf_counter_ns(); source_vector = np.asarray(self.source_model.encode_query(query), dtype="float32"); row["source_query_encode_ms"] = _elapsed(t)
|
|
92
|
+
if source_vector.ndim != 1 or source_vector.shape[0] != int(self.source_model.dimension) or not np.isfinite(source_vector).all():
|
|
93
|
+
raise ValueError("source query encoder returned an invalid vector")
|
|
94
|
+
t = time.perf_counter_ns(); candidates = self.source_index.search(source_vector, K); row["source_ann_ms"] = _elapsed(t)
|
|
95
|
+
candidate_ids = [hit.document_id for hit in candidates]
|
|
96
|
+
t = time.perf_counter_ns(); target_vector = np.asarray(self.target_model.encode_query(query), dtype="float32"); row["target_query_encode_ms"] = _elapsed(t)
|
|
97
|
+
if target_vector.ndim != 1 or target_vector.shape[0] != int(self.target_model.dimension) or not np.isfinite(target_vector).all():
|
|
98
|
+
raise ValueError("target query encoder returned an invalid vector")
|
|
99
|
+
t = time.perf_counter_ns(); cached = self.cache.get(candidate_ids); row["cache_lookup_ms"] = _elapsed(t)
|
|
100
|
+
missing = [doc_id for doc_id in candidate_ids if doc_id not in cached]
|
|
101
|
+
initial_hits = len(cached)
|
|
102
|
+
sync_ids = missing[:max_sync]
|
|
103
|
+
if sync_ids:
|
|
104
|
+
t = time.perf_counter_ns()
|
|
105
|
+
docs = self.documents.get(sync_ids)
|
|
106
|
+
vectors = self.target_model.encode_documents([docs[x] for x in sync_ids], batch_size=self.cfg.migration.background_batch_size)
|
|
107
|
+
self.cache.put(sync_ids, vectors); cached.update({x: np.asarray(v, dtype="float32") for x, v in zip(sync_ids, vectors)})
|
|
108
|
+
row["synchronous_target_encode_ms"] = _elapsed(t)
|
|
109
|
+
else: row["synchronous_target_encode_ms"] = 0.0
|
|
110
|
+
async_ids = [doc_id for doc_id in missing if doc_id not in cached]
|
|
111
|
+
if async_ids: self.worker.enqueue(async_ids)
|
|
112
|
+
scored_ids = [doc_id for doc_id in candidate_ids if doc_id in cached]
|
|
113
|
+
t = time.perf_counter_ns()
|
|
114
|
+
if scored_ids:
|
|
115
|
+
scores = {doc_id: float(np.asarray(cached[doc_id], dtype="float32") @ target_vector) for doc_id in scored_ids}
|
|
116
|
+
candidate_positions = {doc_id: position for position, doc_id in enumerate(candidate_ids)}
|
|
117
|
+
target_order = sorted(scored_ids, key=lambda x: (-scores[x], candidate_positions[x]))
|
|
118
|
+
else:
|
|
119
|
+
scores, target_order = {}, candidate_ids[:]
|
|
120
|
+
row["target_score_ms"] = _elapsed(t)
|
|
121
|
+
t = time.perf_counter_ns(); selected = target_order[:max(1, int(top_k))]; row["topk_ms"] = _elapsed(t)
|
|
122
|
+
target_rank = {doc_id: rank for rank, doc_id in enumerate(target_order, 1)}
|
|
123
|
+
source_rank = {hit.document_id: hit.source_rank + 1 for hit in candidates}
|
|
124
|
+
results = []
|
|
125
|
+
for doc_id in selected:
|
|
126
|
+
results.append({"id": doc_id, "text": self.documents.get([doc_id])[doc_id], "target_score": scores.get(doc_id),
|
|
127
|
+
"source_rank": source_rank.get(doc_id), "target_rank": target_rank.get(doc_id), "target_vector_cached": doc_id in cached})
|
|
128
|
+
# State describes vectors available for this response, including the
|
|
129
|
+
# bounded synchronous budget. ``cache_hits`` remains the persistent
|
|
130
|
+
# lookup count; callers can use ``sync_encoded`` and
|
|
131
|
+
# ``target_vectors_available`` to distinguish a cold lookup that was
|
|
132
|
+
# fully satisfied synchronously from a genuinely partial ranking.
|
|
133
|
+
available_count = len(scored_ids)
|
|
134
|
+
if available_count == len(candidate_ids): status = "WARM"
|
|
135
|
+
elif available_count == 0: status = "COLD"
|
|
136
|
+
else: status = "PARTIAL"
|
|
137
|
+
row["total_ms"] = _elapsed(total_start); row.update({"cache_hits": initial_hits, "cache_misses": len(missing),
|
|
138
|
+
"sync_misses_encoded": len(sync_ids), "async_misses_queued": len(async_ids), "status": status,
|
|
139
|
+
"target_vectors_available": available_count, "candidate_ids": candidate_ids})
|
|
140
|
+
self._record(row)
|
|
141
|
+
timings = {k: float(row[k]) for k in ("source_query_encode_ms", "source_ann_ms", "target_query_encode_ms", "cache_lookup_ms",
|
|
142
|
+
"synchronous_target_encode_ms", "target_score_ms", "topk_ms", "total_ms")}
|
|
143
|
+
return {"results": results, "migration": {"status": status, "source_candidates": len(candidate_ids), "target_cache_hits": initial_hits,
|
|
144
|
+
"target_cache_misses": len(missing), "cache_hit_rate": initial_hits / max(1, len(candidate_ids)),
|
|
145
|
+
"target_vectors_available": available_count,
|
|
146
|
+
"ranking_cache_hit_rate": available_count / max(1, len(candidate_ids)),
|
|
147
|
+
"sync_misses_encoded": len(sync_ids), "async_misses_queued": len(async_ids),
|
|
148
|
+
# Short aliases mirror the public REST schema while the
|
|
149
|
+
# target_* names remain for backwards compatibility.
|
|
150
|
+
"cache_hits": initial_hits, "cache_misses": len(missing),
|
|
151
|
+
"sync_encoded": len(sync_ids), "async_queued": len(async_ids),
|
|
152
|
+
"candidate_depth": K},
|
|
153
|
+
"timing_ms": timings}
|
|
154
|
+
|
|
155
|
+
def prewarm(self, document_ids: list[str], asynchronous: bool = False) -> dict[str, Any]:
|
|
156
|
+
self._ensure_open()
|
|
157
|
+
ids = list(dict.fromkeys(str(x) for x in document_ids)); missing = [x for x in ids if x not in self.cache.contains(ids)]
|
|
158
|
+
if asynchronous:
|
|
159
|
+
queued = self.worker.enqueue(missing); return {"requested": len(ids), "queued": queued, "asynchronous": True}
|
|
160
|
+
if missing:
|
|
161
|
+
batch_size = max(1, int(self.cfg.migration.background_batch_size))
|
|
162
|
+
for start in range(0, len(missing), batch_size):
|
|
163
|
+
chunk = missing[start:start + batch_size]
|
|
164
|
+
docs = self.documents.get(chunk)
|
|
165
|
+
vectors = self.target_model.encode_documents([docs[x] for x in chunk], batch_size=batch_size)
|
|
166
|
+
self.cache.put(chunk, vectors)
|
|
167
|
+
return {"requested": len(ids), "materialized": len(missing), "asynchronous": False}
|
|
168
|
+
|
|
169
|
+
def status(self) -> dict[str, Any]:
|
|
170
|
+
self._ensure_open()
|
|
171
|
+
cache = self.cache.stats(); worker = self.worker.stats(); data = self.state.snapshot()
|
|
172
|
+
plan = make_plan(target_model=self.cfg.target.model, corpus_size=self.documents.size(), diagnostic=self.plan.diagnostic,
|
|
173
|
+
recommended_k=self.plan.candidate_depth, ann_status=self.ann_status,
|
|
174
|
+
cached_target_vectors=cache["cached_target_vectors"], default_k=self.cfg.migration.candidate_depth,
|
|
175
|
+
source_model=self.cfg.source.model)
|
|
176
|
+
values: list[float] = []
|
|
177
|
+
for row in self.records:
|
|
178
|
+
try:
|
|
179
|
+
value = float(row.get("total_ms"))
|
|
180
|
+
except (TypeError, ValueError):
|
|
181
|
+
continue
|
|
182
|
+
if np.isfinite(value) and value >= 0:
|
|
183
|
+
values.append(value)
|
|
184
|
+
lat = summarize(values) if values else {"count": 0}
|
|
185
|
+
return {"migration": plan.to_dict(), "cache": cache, "worker": worker, "state": data, "latency": lat}
|
|
186
|
+
|
|
187
|
+
def close(self, close_models: bool = True, close_indexes: bool = True) -> None:
|
|
188
|
+
"""Stop background work and close resources.
|
|
189
|
+
|
|
190
|
+
``close_models=False`` is used by the public facade when the caller
|
|
191
|
+
supplied model instances that it still owns. Existing CLI callers
|
|
192
|
+
retain the original behavior through the default.
|
|
193
|
+
"""
|
|
194
|
+
if self._closed:
|
|
195
|
+
return
|
|
196
|
+
try:
|
|
197
|
+
self.worker.close()
|
|
198
|
+
finally:
|
|
199
|
+
try:
|
|
200
|
+
self.cache.close()
|
|
201
|
+
finally:
|
|
202
|
+
try:
|
|
203
|
+
# Qdrant owns a file lock/client connection; release both
|
|
204
|
+
# source and optional target indexes when the engine owns
|
|
205
|
+
# the serving session. In-memory/FAISS implementations are
|
|
206
|
+
# no-ops, and identity avoids closing a shared object twice.
|
|
207
|
+
if close_indexes:
|
|
208
|
+
closed_indexes: set[int] = set()
|
|
209
|
+
for index in (self.source_index, self.target_index):
|
|
210
|
+
if index is None or id(index) in closed_indexes:
|
|
211
|
+
continue
|
|
212
|
+
close_index = getattr(index, "close", None)
|
|
213
|
+
if callable(close_index):
|
|
214
|
+
close_index()
|
|
215
|
+
closed_indexes.add(id(index))
|
|
216
|
+
finally:
|
|
217
|
+
if close_models:
|
|
218
|
+
try:
|
|
219
|
+
self.source_model.close()
|
|
220
|
+
finally:
|
|
221
|
+
self.target_model.close()
|
|
222
|
+
self._closed = True
|
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
"""Pydantic request/response models for the public REST API."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from typing import Any
|
|
6
|
+
|
|
7
|
+
try:
|
|
8
|
+
from pydantic import BaseModel, Field
|
|
9
|
+
except ImportError: # pragma: no cover - imported only when API extras are used
|
|
10
|
+
BaseModel = object # type: ignore[misc,assignment]
|
|
11
|
+
|
|
12
|
+
def Field(default: Any = None, **_: Any) -> Any: # type: ignore[misc]
|
|
13
|
+
return default
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
class SearchRequest(BaseModel):
|
|
17
|
+
query: str = Field(min_length=1)
|
|
18
|
+
top_k: int = Field(default=10, ge=1, le=100)
|
|
19
|
+
candidate_depth: int | None = Field(default=None, ge=1)
|
|
20
|
+
max_sync_misses: int | None = Field(default=None, ge=0)
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class PrewarmRequest(BaseModel):
|
|
24
|
+
document_ids: list[str] = Field(default_factory=list)
|
|
25
|
+
asynchronous: bool = True
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
class SearchResponse(BaseModel):
|
|
29
|
+
results: list[dict[str, Any]]
|
|
30
|
+
migration: dict[str, Any]
|
|
31
|
+
timing_ms: dict[str, float]
|
|
32
|
+
# Flattened aliases make the response convenient for clients while the
|
|
33
|
+
# nested ``migration`` object remains for backwards compatibility.
|
|
34
|
+
state: str | None = None
|
|
35
|
+
candidate_depth: int | None = None
|
|
36
|
+
cache_hits: int | None = None
|
|
37
|
+
cache_misses: int | None = None
|
|
38
|
+
sync_encoded: int | None = None
|
|
39
|
+
async_queued: int | None = None
|
|
@@ -0,0 +1,210 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: embedflow
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Progressive embedding-model migration over existing vector indexes.
|
|
5
|
+
Author: Arnav Srivastav
|
|
6
|
+
License-Expression: Apache-2.0
|
|
7
|
+
Project-URL: Homepage, https://embedflow.org
|
|
8
|
+
Project-URL: Repository, https://github.com/arnsri33/embedflow
|
|
9
|
+
Project-URL: Documentation, https://github.com/arnsri33/embedflow#readme
|
|
10
|
+
Project-URL: Issues, https://github.com/arnsri33/embedflow/issues
|
|
11
|
+
Keywords: embeddings,vector-search,rag,information-retrieval,faiss,qdrant
|
|
12
|
+
Classifier: Development Status :: 3 - Alpha
|
|
13
|
+
Classifier: Intended Audience :: Developers
|
|
14
|
+
Classifier: Intended Audience :: Science/Research
|
|
15
|
+
Classifier: Programming Language :: Python :: 3
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
19
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
20
|
+
Requires-Python: >=3.10
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
License-File: LICENSE
|
|
23
|
+
Requires-Dist: numpy>=1.24
|
|
24
|
+
Requires-Dist: pyyaml>=6.0
|
|
25
|
+
Requires-Dist: tqdm>=4.65
|
|
26
|
+
Provides-Extra: faiss
|
|
27
|
+
Requires-Dist: faiss-cpu>=1.8.0; extra == "faiss"
|
|
28
|
+
Provides-Extra: qdrant
|
|
29
|
+
Requires-Dist: qdrant-client>=1.9; extra == "qdrant"
|
|
30
|
+
Provides-Extra: models
|
|
31
|
+
Requires-Dist: huggingface-hub>=0.20; extra == "models"
|
|
32
|
+
Requires-Dist: transformers>=4.45; extra == "models"
|
|
33
|
+
Requires-Dist: sentence-transformers>=3.0; extra == "models"
|
|
34
|
+
Requires-Dist: torch>=2.2; extra == "models"
|
|
35
|
+
Provides-Extra: dashboard
|
|
36
|
+
Requires-Dist: fastapi>=0.100; extra == "dashboard"
|
|
37
|
+
Requires-Dist: uvicorn[standard]>=0.23; extra == "dashboard"
|
|
38
|
+
Requires-Dist: pydantic>=2.0; extra == "dashboard"
|
|
39
|
+
Provides-Extra: all
|
|
40
|
+
Requires-Dist: faiss-cpu>=1.8.0; extra == "all"
|
|
41
|
+
Requires-Dist: qdrant-client>=1.9; extra == "all"
|
|
42
|
+
Requires-Dist: huggingface-hub>=0.20; extra == "all"
|
|
43
|
+
Requires-Dist: transformers>=4.45; extra == "all"
|
|
44
|
+
Requires-Dist: sentence-transformers>=3.0; extra == "all"
|
|
45
|
+
Requires-Dist: torch>=2.2; extra == "all"
|
|
46
|
+
Requires-Dist: fastapi>=0.100; extra == "all"
|
|
47
|
+
Requires-Dist: uvicorn[standard]>=0.23; extra == "all"
|
|
48
|
+
Requires-Dist: pydantic>=2.0; extra == "all"
|
|
49
|
+
Provides-Extra: dev
|
|
50
|
+
Requires-Dist: pytest>=8.0; extra == "dev"
|
|
51
|
+
Requires-Dist: ruff>=0.5; extra == "dev"
|
|
52
|
+
Requires-Dist: build>=1.2; extra == "dev"
|
|
53
|
+
Dynamic: license-file
|
|
54
|
+
|
|
55
|
+
# EmbedFlow
|
|
56
|
+
|
|
57
|
+
**Progressive embedding-model migration over existing vector indexes.**
|
|
58
|
+
|
|
59
|
+
EmbedFlow lets a new embedding model serve over candidates from an existing
|
|
60
|
+
vector index while target document vectors are materialized progressively. It
|
|
61
|
+
supports migration analysis, persistent caching, background work, FAISS,
|
|
62
|
+
Qdrant, a CLI, and FastAPI.
|
|
63
|
+
|
|
64
|
+
The full project README and architecture diagram are on
|
|
65
|
+
<https://github.com/arnsri33/embedflow>.
|
|
66
|
+
|
|
67
|
+
## Install
|
|
68
|
+
|
|
69
|
+
```bash
|
|
70
|
+
python -m pip install embedflow
|
|
71
|
+
```
|
|
72
|
+
|
|
73
|
+
For FAISS and the dashboard:
|
|
74
|
+
|
|
75
|
+
```bash
|
|
76
|
+
python -m pip install "embedflow[faiss,dashboard]"
|
|
77
|
+
```
|
|
78
|
+
|
|
79
|
+
Qdrant and model-runtime extras are documented in the
|
|
80
|
+
[installation guide](https://github.com/arnsri33/embedflow/blob/main/docs/installation.md).
|
|
81
|
+
For model-backed analysis, install `embedflow[faiss,models,dashboard]`.
|
|
82
|
+
|
|
83
|
+
## Why EmbedFlow?
|
|
84
|
+
|
|
85
|
+
Embedding-model upgrades usually mean re-embedding the corpus and building a
|
|
86
|
+
second index before the new model can serve. EmbedFlow tests whether the
|
|
87
|
+
existing retriever can remain useful during that transition.
|
|
88
|
+
|
|
89
|
+
Different representation spaces can still preserve useful retrieval
|
|
90
|
+
neighborhoods. EmbedFlow retrieves a bounded source candidate set, scores those
|
|
91
|
+
candidates with the target model, and fills a persistent target-vector cache in
|
|
92
|
+
the background.
|
|
93
|
+
|
|
94
|
+
## Try it
|
|
95
|
+
|
|
96
|
+
The deterministic demo needs no paid service or model download. From a source
|
|
97
|
+
checkout, run:
|
|
98
|
+
|
|
99
|
+
```bash
|
|
100
|
+
./scripts/run_demo.sh
|
|
101
|
+
```
|
|
102
|
+
|
|
103
|
+
Open <http://127.0.0.1:8000/>. The first search can be `COLD` or `PARTIAL`;
|
|
104
|
+
repeated traffic becomes `WARM` as target vectors are materialized.
|
|
105
|
+
|
|
106
|
+
## Analyze a migration
|
|
107
|
+
|
|
108
|
+
```bash
|
|
109
|
+
embedflow analyze \
|
|
110
|
+
--documents ./documents.jsonl \
|
|
111
|
+
--index ./legacy.index \
|
|
112
|
+
--source-model sentence-transformers/all-MiniLM-L6-v2 \
|
|
113
|
+
--target-model Qwen/Qwen3-Embedding-0.6B \
|
|
114
|
+
--probe-queries ./probe_queries.jsonl \
|
|
115
|
+
--device cuda \
|
|
116
|
+
--output-dir ./analysis
|
|
117
|
+
```
|
|
118
|
+
|
|
119
|
+
The report includes the frozen T2-v1 diagnostic (`SAFE`, `EXPAND`, or
|
|
120
|
+
`UNSAFE_OR_UNCERTAIN`), a recommended initial candidate depth, and ANN health.
|
|
121
|
+
`SAFE` is an empirical deployment signal; validate important migrations on the
|
|
122
|
+
target corpus.
|
|
123
|
+
|
|
124
|
+
Start progressive serving with:
|
|
125
|
+
|
|
126
|
+
```bash
|
|
127
|
+
embedflow serve --config ./analysis/embedflow.analysis.yaml --device cuda
|
|
128
|
+
```
|
|
129
|
+
|
|
130
|
+
## Known migration evidence
|
|
131
|
+
|
|
132
|
+
EmbedFlow ships a versioned registry of measured results from the research
|
|
133
|
+
study. Matching model contracts provide useful starting depths and show what
|
|
134
|
+
was observed on earlier corpora; a new corpus still receives its own analysis.
|
|
135
|
+
|
|
136
|
+
```bash
|
|
137
|
+
embedflow registry list
|
|
138
|
+
embedflow registry show \
|
|
139
|
+
--source Qwen/Qwen3-Embedding-4B \
|
|
140
|
+
--target Qwen/Qwen3-Embedding-8B
|
|
141
|
+
```
|
|
142
|
+
|
|
143
|
+
Selected core records (nDCG@10, `G(50)`):
|
|
144
|
+
|
|
145
|
+
| Source | Target | Evaluation | `G(50)` | Observed depth |
|
|
146
|
+
| --- | --- | --- | ---: | ---: |
|
|
147
|
+
| MiniLM-L6-v2 | Qwen3-8B | BRIGHT, 413K | 0.03665 | — |
|
|
148
|
+
| Qwen3-0.6B | Qwen3-8B | BRIGHT, 413K | 0.01465 | 200 |
|
|
149
|
+
| Qwen3-4B | Qwen3-8B | BRIGHT, 413K | 0.00347 | 20 |
|
|
150
|
+
| MiniLM-L6-v2 | Qwen3-8B | Natural Questions, 1M | 0.03255 | 500 |
|
|
151
|
+
| Qwen3-4B | Qwen3-8B | Natural Questions, 1M | -0.00043 | 20 |
|
|
152
|
+
|
|
153
|
+
See the [registry documentation](https://github.com/arnsri33/embedflow/blob/main/docs/registry.md)
|
|
154
|
+
for matching levels, contract fingerprints, and provenance.
|
|
155
|
+
|
|
156
|
+
## Research
|
|
157
|
+
|
|
158
|
+
For candidate depth `K`, EmbedFlow measures:
|
|
159
|
+
|
|
160
|
+
```text
|
|
161
|
+
G(K) = M_T - M_{T|S_K}
|
|
162
|
+
```
|
|
163
|
+
|
|
164
|
+
Lower `G(K)` means the source candidate pool recovers more of native target
|
|
165
|
+
retrieval quality. Containment reports neighborhood overlap separately. The
|
|
166
|
+
study includes 63 development settings, frozen BRIGHT validation, and Natural
|
|
167
|
+
Questions scale experiments through 1M documents.
|
|
168
|
+
|
|
169
|
+
Read the [concepts](https://github.com/arnsri33/embedflow/blob/main/docs/concepts.md)
|
|
170
|
+
and [methodology](https://github.com/arnsri33/embedflow/blob/main/docs/methodology.md)
|
|
171
|
+
for definitions and reproduction details.
|
|
172
|
+
|
|
173
|
+
## Integrations
|
|
174
|
+
|
|
175
|
+
| Backend | Status |
|
|
176
|
+
| --- | --- |
|
|
177
|
+
| FAISS | Supported |
|
|
178
|
+
| Qdrant | Supported |
|
|
179
|
+
|
|
180
|
+
See the [FAISS guide](https://github.com/arnsri33/embedflow/blob/main/docs/integrations/faiss.md)
|
|
181
|
+
and [Qdrant guide](https://github.com/arnsri33/embedflow/blob/main/docs/integrations/qdrant.md).
|
|
182
|
+
|
|
183
|
+
## CLI
|
|
184
|
+
|
|
185
|
+
```bash
|
|
186
|
+
embedflow --help
|
|
187
|
+
embedflow analyze --help
|
|
188
|
+
embedflow serve --config ./embedflow.yaml
|
|
189
|
+
embedflow status --config ./embedflow.yaml
|
|
190
|
+
embedflow registry list
|
|
191
|
+
embedflow economics --corpus-size 1000000000 --docs-per-second 106.98 --gpu-price 3.29
|
|
192
|
+
embedflow doctor --config ./embedflow.yaml
|
|
193
|
+
```
|
|
194
|
+
|
|
195
|
+
The [CLI reference](https://github.com/arnsri33/embedflow/blob/main/docs/cli.md)
|
|
196
|
+
and [API reference](https://github.com/arnsri33/embedflow/blob/main/docs/api.md)
|
|
197
|
+
cover the remaining commands and endpoints.
|
|
198
|
+
|
|
199
|
+
## Status
|
|
200
|
+
|
|
201
|
+
EmbedFlow v0.1.0 is an alpha release for research and early real-world
|
|
202
|
+
testing. T2-v1 is an empirical finite-tail diagnostic, partial rankings can
|
|
203
|
+
differ from fully warm target reranking, and ANN fidelity needs a reference
|
|
204
|
+
comparison to audit.
|
|
205
|
+
|
|
206
|
+
## Links
|
|
207
|
+
|
|
208
|
+
- [GitHub repository](https://github.com/arnsri33/embedflow)
|
|
209
|
+
- [Documentation](https://github.com/arnsri33/embedflow#readme)
|
|
210
|
+
- [License](https://github.com/arnsri33/embedflow/blob/main/LICENSE)
|