|
|
@@ -0,0 +1,591 @@
|
|
|
+"""Build the local Script Decode vector index from immutable JSON inputs.
|
|
|
+
|
|
|
+The builder is intentionally self-contained in this repository. It never imports
|
|
|
+the previous-generation application and never mutates its raw JSON inputs.
|
|
|
+"""
|
|
|
+
|
|
|
+from __future__ import annotations
|
|
|
+
|
|
|
+import argparse
|
|
|
+import asyncio
|
|
|
+import json
|
|
|
+import os
|
|
|
+import re
|
|
|
+import shutil
|
|
|
+import time
|
|
|
+from collections.abc import Iterator, Mapping, Sequence
|
|
|
+from dataclasses import dataclass
|
|
|
+from datetime import UTC, datetime
|
|
|
+from hashlib import sha256
|
|
|
+from pathlib import Path
|
|
|
+from typing import Any
|
|
|
+from urllib.parse import urlsplit
|
|
|
+from uuid import uuid4
|
|
|
+
|
|
|
+import httpx
|
|
|
+import numpy as np
|
|
|
+
|
|
|
+_COLUMN_KEY_MAP = {
|
|
|
+ "主题原子点": "主题",
|
|
|
+ "形式原子点": "形式",
|
|
|
+ "作用原子点": "作用",
|
|
|
+ "感受原子点": "感受",
|
|
|
+}
|
|
|
+_TEXT_FIELD_TO_SUB = {
|
|
|
+ "维度名": "维度名",
|
|
|
+ "维度名描述": "维度名",
|
|
|
+ "维度值": "维度值",
|
|
|
+ "维度值描述": "维度值",
|
|
|
+ "维度": "维度名",
|
|
|
+ "原子点": "维度值",
|
|
|
+}
|
|
|
+_RETRYABLE_STATUS_CODES = frozenset({408, 409, 425, 429, 500, 502, 503, 504})
|
|
|
+
|
|
|
+
|
|
|
+class EmbeddingApiError(RuntimeError):
|
|
|
+ """The remote embedding API failed or returned an invalid contract."""
|
|
|
+
|
|
|
+
|
|
|
+@dataclass(frozen=True, slots=True)
|
|
|
+class DecodeIndexBuildResult:
|
|
|
+ index_root: Path
|
|
|
+ manifest: dict[str, Any]
|
|
|
+
|
|
|
+
|
|
|
+class DashScopeEmbeddingClient:
|
|
|
+ """Bounded async client for DashScope's OpenAI-compatible embedding API."""
|
|
|
+
|
|
|
+ def __init__(
|
|
|
+ self,
|
|
|
+ http: httpx.AsyncClient,
|
|
|
+ *,
|
|
|
+ api_key: str,
|
|
|
+ endpoint: str,
|
|
|
+ model: str,
|
|
|
+ dimension: int,
|
|
|
+ concurrency: int = 8,
|
|
|
+ max_retries: int = 6,
|
|
|
+ ) -> None:
|
|
|
+ if not api_key:
|
|
|
+ raise ValueError("DashScope API key is required")
|
|
|
+ if dimension < 1:
|
|
|
+ raise ValueError("embedding dimension must be positive")
|
|
|
+ if concurrency < 1:
|
|
|
+ raise ValueError("concurrency must be positive")
|
|
|
+ if max_retries < 0:
|
|
|
+ raise ValueError("max_retries cannot be negative")
|
|
|
+ self._http = http
|
|
|
+ self._headers = {
|
|
|
+ "Authorization": f"Bearer {api_key}",
|
|
|
+ "Content-Type": "application/json",
|
|
|
+ }
|
|
|
+ self.endpoint = endpoint
|
|
|
+ self.model = model
|
|
|
+ self.dimension = dimension
|
|
|
+ self._semaphore = asyncio.Semaphore(concurrency)
|
|
|
+ self._max_retries = max_retries
|
|
|
+
|
|
|
+ async def embed(self, texts: Sequence[str]) -> tuple[Any, int]:
|
|
|
+ if not texts or len(texts) > 10:
|
|
|
+ raise ValueError("DashScope synchronous embedding batches must contain 1-10 texts")
|
|
|
+ async with self._semaphore:
|
|
|
+ return await self._request(texts)
|
|
|
+
|
|
|
+ async def _request(self, texts: Sequence[str]) -> tuple[Any, int]:
|
|
|
+ last_error: Exception | None = None
|
|
|
+ for attempt in range(self._max_retries + 1):
|
|
|
+ try:
|
|
|
+ response = await self._http.post(
|
|
|
+ self.endpoint,
|
|
|
+ headers=self._headers,
|
|
|
+ json={
|
|
|
+ "model": self.model,
|
|
|
+ "input": list(texts),
|
|
|
+ "encoding_format": "float",
|
|
|
+ "dimensions": self.dimension,
|
|
|
+ },
|
|
|
+ )
|
|
|
+ except httpx.HTTPError as exc:
|
|
|
+ last_error = exc
|
|
|
+ if attempt >= self._max_retries:
|
|
|
+ break
|
|
|
+ await asyncio.sleep(min(2**attempt, 16))
|
|
|
+ continue
|
|
|
+ if response.status_code in _RETRYABLE_STATUS_CODES and attempt < self._max_retries:
|
|
|
+ retry_after = _retry_after_seconds(response)
|
|
|
+ await asyncio.sleep(retry_after if retry_after is not None else min(2**attempt, 16))
|
|
|
+ continue
|
|
|
+ if not response.is_success:
|
|
|
+ raise EmbeddingApiError(_bounded_api_error(response))
|
|
|
+ return _parse_embedding_response(response, len(texts), self.dimension)
|
|
|
+ raise EmbeddingApiError("embedding request failed after bounded retries") from last_error
|
|
|
+
|
|
|
+
|
|
|
+def extract_decode_chunks(raw_root: Path) -> list[dict[str, Any]]:
|
|
|
+ """Extract the exact searchable fields used by the previous Script Build flow."""
|
|
|
+
|
|
|
+ resolved = raw_root.resolve()
|
|
|
+ if not resolved.is_dir():
|
|
|
+ raise FileNotFoundError(f"decode raw-data directory does not exist: {resolved}")
|
|
|
+ files = sorted(resolved.rglob("*.json"))
|
|
|
+ if not files:
|
|
|
+ raise RuntimeError(f"decode raw-data directory has no JSON files: {resolved}")
|
|
|
+ chunks: list[dict[str, Any]] = []
|
|
|
+ chunk_ids: set[str] = set()
|
|
|
+ for path in files:
|
|
|
+ try:
|
|
|
+ payload = json.loads(path.read_text(encoding="utf-8"))
|
|
|
+ except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
|
|
|
+ raise RuntimeError(f"decode source is not valid UTF-8 JSON: {path}") from exc
|
|
|
+ if not isinstance(payload, dict):
|
|
|
+ raise RuntimeError(f"decode source root must be an object: {path}")
|
|
|
+ account = path.parent.name
|
|
|
+ post_id = str(payload.get("帖子ID") or "").strip() or path.stem
|
|
|
+ relpath = (Path(resolved.name) / path.relative_to(resolved)).as_posix()
|
|
|
+ for paragraph in _iter_paragraphs(payload.get("段落列表")):
|
|
|
+ paragraph_id = str(paragraph.get("id") or "")
|
|
|
+ paragraph_name = str(paragraph.get("名称") or "")
|
|
|
+ for atom_key, column in _COLUMN_KEY_MAP.items():
|
|
|
+ atoms = paragraph.get(atom_key)
|
|
|
+ if not isinstance(atoms, list):
|
|
|
+ continue
|
|
|
+ for atom_index, atom in enumerate(atoms):
|
|
|
+ if not isinstance(atom, dict):
|
|
|
+ continue
|
|
|
+ for text_field, sub_field in _TEXT_FIELD_TO_SUB.items():
|
|
|
+ text = str(atom.get(text_field) or "").strip()
|
|
|
+ if not text:
|
|
|
+ continue
|
|
|
+ chunk_id = f"{post_id}|{paragraph_id}|{column}|{atom_index}|{text_field}"
|
|
|
+ if chunk_id in chunk_ids:
|
|
|
+ raise RuntimeError(f"duplicate decode chunk identity: {chunk_id}")
|
|
|
+ chunk_ids.add(chunk_id)
|
|
|
+ chunks.append(
|
|
|
+ {
|
|
|
+ "chunk_id": chunk_id,
|
|
|
+ "post_id": post_id,
|
|
|
+ "account": account,
|
|
|
+ "json_relpath": relpath,
|
|
|
+ "column": column,
|
|
|
+ "sub_field": sub_field,
|
|
|
+ "text_field": text_field,
|
|
|
+ "text": text,
|
|
|
+ "paragraph_id": paragraph_id,
|
|
|
+ "paragraph_name": paragraph_name,
|
|
|
+ "atom_index": atom_index,
|
|
|
+ }
|
|
|
+ )
|
|
|
+ if not chunks:
|
|
|
+ raise RuntimeError("decode sources produced no searchable chunks")
|
|
|
+ return chunks
|
|
|
+
|
|
|
+
|
|
|
+async def rebuild_decode_index(
|
|
|
+ *,
|
|
|
+ raw_root: Path,
|
|
|
+ index_root: Path,
|
|
|
+ client: DashScopeEmbeddingClient,
|
|
|
+ batch_size: int = 10,
|
|
|
+ shard_size: int = 1_000,
|
|
|
+ storage_dtype: str = "float16",
|
|
|
+ reset_work: bool = False,
|
|
|
+) -> DecodeIndexBuildResult:
|
|
|
+ """Fully rebuild and atomically replace an index, retaining resumable work on failure."""
|
|
|
+
|
|
|
+ if not 1 <= batch_size <= 10:
|
|
|
+ raise ValueError("batch_size must be between 1 and 10 for DashScope")
|
|
|
+ if shard_size < batch_size:
|
|
|
+ raise ValueError("shard_size must be at least one batch")
|
|
|
+ if storage_dtype not in {"float16", "float32"}:
|
|
|
+ raise ValueError("storage_dtype must be float16 or float32")
|
|
|
+ raw_root, index_root = await asyncio.gather(
|
|
|
+ asyncio.to_thread(raw_root.resolve),
|
|
|
+ asyncio.to_thread(index_root.resolve),
|
|
|
+ )
|
|
|
+ if raw_root == index_root or index_root == index_root.parent:
|
|
|
+ raise ValueError("raw_root and index_root must be separate bounded directories")
|
|
|
+
|
|
|
+ started = time.monotonic()
|
|
|
+ chunks = await asyncio.to_thread(extract_decode_chunks, raw_root)
|
|
|
+ source_digest = await asyncio.to_thread(_directory_digest, raw_root)
|
|
|
+ chunks_bytes = _json_lines(chunks)
|
|
|
+ chunks_digest = "sha256:" + sha256(chunks_bytes).hexdigest()
|
|
|
+ fingerprint_payload = {
|
|
|
+ "builder": "script-build-decode-index/v1",
|
|
|
+ "source_sha256": source_digest,
|
|
|
+ "meta_sha256": chunks_digest,
|
|
|
+ "model": client.model,
|
|
|
+ "dimension": client.dimension,
|
|
|
+ "storage_dtype": storage_dtype,
|
|
|
+ "batch_size": batch_size,
|
|
|
+ "shard_size": shard_size,
|
|
|
+ }
|
|
|
+ fingerprint = sha256(
|
|
|
+ json.dumps(fingerprint_payload, ensure_ascii=False, sort_keys=True).encode("utf-8")
|
|
|
+ ).hexdigest()
|
|
|
+ safe_model = re.sub(r"[^A-Za-z0-9._-]+", "-", client.model).strip("-") or "model"
|
|
|
+ work_root = index_root.parent / (f".{index_root.name}-build-{safe_model}-{client.dimension}")
|
|
|
+ if reset_work and work_root.exists():
|
|
|
+ await asyncio.to_thread(shutil.rmtree, work_root)
|
|
|
+ work_root.mkdir(parents=True, exist_ok=True)
|
|
|
+ state_path = work_root / "state.json"
|
|
|
+ if state_path.exists():
|
|
|
+ state = _read_json_object(state_path)
|
|
|
+ if state.get("fingerprint") != fingerprint:
|
|
|
+ raise RuntimeError(
|
|
|
+ f"resumable work belongs to different inputs/configuration: {work_root}; "
|
|
|
+ "rerun with --reset-work"
|
|
|
+ )
|
|
|
+ else:
|
|
|
+ _atomic_write_json(
|
|
|
+ state_path,
|
|
|
+ {"fingerprint": fingerprint, "configuration": fingerprint_payload},
|
|
|
+ )
|
|
|
+ meta_work = work_root / "meta.jsonl"
|
|
|
+ if meta_work.exists() and _file_digest(meta_work) != chunks_digest:
|
|
|
+ raise RuntimeError(f"resumable metadata is inconsistent: {meta_work}")
|
|
|
+ if not meta_work.exists():
|
|
|
+ _atomic_write_bytes(meta_work, chunks_bytes)
|
|
|
+
|
|
|
+ shard_root = work_root / "shards"
|
|
|
+ shard_root.mkdir(exist_ok=True)
|
|
|
+ total_tokens = 0
|
|
|
+ request_count = 0
|
|
|
+ shard_paths: list[Path] = []
|
|
|
+ for start in range(0, len(chunks), shard_size):
|
|
|
+ end = min(start + shard_size, len(chunks))
|
|
|
+ shard_path = shard_root / f"{start:06d}-{end:06d}.npy"
|
|
|
+ usage_path = shard_path.with_suffix(".json")
|
|
|
+ if shard_path.exists() and usage_path.exists():
|
|
|
+ matrix = await asyncio.to_thread(np.load, shard_path, mmap_mode="r")
|
|
|
+ if matrix.shape != (end - start, client.dimension):
|
|
|
+ raise RuntimeError(f"resumable vector shard has an invalid shape: {shard_path}")
|
|
|
+ usage = _read_json_object(usage_path)
|
|
|
+ total_tokens += int(usage.get("prompt_tokens", 0))
|
|
|
+ request_count += int(usage.get("request_count", 0))
|
|
|
+ shard_paths.append(shard_path)
|
|
|
+ print(f"↪ 已恢复 {end}/{len(chunks)} 条向量", flush=True)
|
|
|
+ continue
|
|
|
+ texts = [str(chunk["text"]) for chunk in chunks[start:end]]
|
|
|
+ matrix, shard_tokens, shard_requests = await _embed_texts(
|
|
|
+ client, texts, batch_size=batch_size
|
|
|
+ )
|
|
|
+ target_dtype = np.float16 if storage_dtype == "float16" else np.float32
|
|
|
+ matrix = np.asarray(matrix, dtype=target_dtype)
|
|
|
+ _atomic_save_array(shard_path, matrix)
|
|
|
+ _atomic_write_json(
|
|
|
+ usage_path,
|
|
|
+ {
|
|
|
+ "prompt_tokens": shard_tokens,
|
|
|
+ "request_count": shard_requests,
|
|
|
+ "rows": end - start,
|
|
|
+ },
|
|
|
+ )
|
|
|
+ total_tokens += shard_tokens
|
|
|
+ request_count += shard_requests
|
|
|
+ shard_paths.append(shard_path)
|
|
|
+ print(
|
|
|
+ f"✓ 已生成 {end}/{len(chunks)} 条向量, 累计 tokens={total_tokens}",
|
|
|
+ flush=True,
|
|
|
+ )
|
|
|
+
|
|
|
+ stage_root = index_root.parent / f".{index_root.name}-publish-{uuid4().hex}"
|
|
|
+ stage_root.mkdir(parents=True)
|
|
|
+ try:
|
|
|
+ shutil.copyfile(meta_work, stage_root / "meta.jsonl")
|
|
|
+ target_dtype = np.float16 if storage_dtype == "float16" else np.float32
|
|
|
+ final_matrix = np.lib.format.open_memmap(
|
|
|
+ stage_root / "embeddings.npy",
|
|
|
+ mode="w+",
|
|
|
+ dtype=target_dtype,
|
|
|
+ shape=(len(chunks), client.dimension),
|
|
|
+ )
|
|
|
+ offset = 0
|
|
|
+ for shard_path in shard_paths:
|
|
|
+ shard = np.load(shard_path, mmap_mode="r")
|
|
|
+ final_matrix[offset : offset + len(shard)] = shard
|
|
|
+ offset += len(shard)
|
|
|
+ final_matrix.flush()
|
|
|
+ del final_matrix
|
|
|
+ verification = await asyncio.to_thread(
|
|
|
+ _verify_vectors,
|
|
|
+ stage_root / "embeddings.npy",
|
|
|
+ len(chunks),
|
|
|
+ client.dimension,
|
|
|
+ )
|
|
|
+ embeddings_digest = await asyncio.to_thread(_file_digest, stage_root / "embeddings.npy")
|
|
|
+ endpoint_host = urlsplit(client.endpoint).hostname or ""
|
|
|
+ manifest: dict[str, Any] = {
|
|
|
+ "built_at": datetime.now(UTC).isoformat(),
|
|
|
+ "builder": "script-build-decode-index/v1",
|
|
|
+ "provider": "aliyun-bailian",
|
|
|
+ "endpoint_host": endpoint_host,
|
|
|
+ "model": client.model,
|
|
|
+ "dimension": client.dimension,
|
|
|
+ "storage_dtype": storage_dtype,
|
|
|
+ "post_count": len({str(chunk["post_id"]) for chunk in chunks}),
|
|
|
+ "account_count": len({str(chunk["account"]) for chunk in chunks}),
|
|
|
+ "chunk_count": len(chunks),
|
|
|
+ "total_tokens": total_tokens,
|
|
|
+ "request_count": request_count,
|
|
|
+ "account_filter": None,
|
|
|
+ "elapsed_seconds": round(time.monotonic() - started, 2),
|
|
|
+ "embeddings_file_mb": round(
|
|
|
+ (stage_root / "embeddings.npy").stat().st_size / (1024 * 1024), 2
|
|
|
+ ),
|
|
|
+ "source_file_count": len(list(raw_root.rglob("*.json"))),
|
|
|
+ "source_sha256": source_digest,
|
|
|
+ "meta_sha256": chunks_digest,
|
|
|
+ "embeddings_sha256": embeddings_digest,
|
|
|
+ "min_vector_norm": verification[0],
|
|
|
+ "max_vector_norm": verification[1],
|
|
|
+ }
|
|
|
+ _atomic_write_json(stage_root / "manifest.json", manifest)
|
|
|
+ await asyncio.to_thread(_publish_index, stage_root, index_root)
|
|
|
+ except BaseException:
|
|
|
+ if stage_root.exists():
|
|
|
+ shutil.rmtree(stage_root)
|
|
|
+ raise
|
|
|
+ await asyncio.to_thread(shutil.rmtree, work_root)
|
|
|
+ return DecodeIndexBuildResult(index_root=index_root, manifest=manifest)
|
|
|
+
|
|
|
+
|
|
|
+async def _embed_texts(
|
|
|
+ client: DashScopeEmbeddingClient,
|
|
|
+ texts: Sequence[str],
|
|
|
+ *,
|
|
|
+ batch_size: int,
|
|
|
+) -> tuple[Any, int, int]:
|
|
|
+ batches = [texts[index : index + batch_size] for index in range(0, len(texts), batch_size)]
|
|
|
+
|
|
|
+ async def run(index: int, batch: Sequence[str]) -> tuple[int, Any, int]:
|
|
|
+ matrix, tokens = await client.embed(batch)
|
|
|
+ return index, matrix, tokens
|
|
|
+
|
|
|
+ results = await asyncio.gather(*(run(index, batch) for index, batch in enumerate(batches)))
|
|
|
+ ordered = sorted(results, key=lambda result: result[0])
|
|
|
+ return (
|
|
|
+ np.concatenate([result[1] for result in ordered]),
|
|
|
+ sum(result[2] for result in ordered),
|
|
|
+ len(batches),
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+def _iter_paragraphs(value: object) -> Iterator[Mapping[str, Any]]:
|
|
|
+ if not isinstance(value, list):
|
|
|
+ return
|
|
|
+ for paragraph in value:
|
|
|
+ if not isinstance(paragraph, dict):
|
|
|
+ continue
|
|
|
+ yield paragraph
|
|
|
+ yield from _iter_paragraphs(paragraph.get("子项"))
|
|
|
+
|
|
|
+
|
|
|
+def _parse_embedding_response(
|
|
|
+ response: httpx.Response, expected_rows: int, dimension: int
|
|
|
+) -> tuple[Any, int]:
|
|
|
+ try:
|
|
|
+ payload = response.json()
|
|
|
+ except ValueError as exc:
|
|
|
+ raise EmbeddingApiError("embedding response is not JSON") from exc
|
|
|
+ rows = payload.get("data") if isinstance(payload, dict) else None
|
|
|
+ if not isinstance(rows, list) or len(rows) != expected_rows:
|
|
|
+ raise EmbeddingApiError("embedding response row count does not match the request")
|
|
|
+ if all(isinstance(row, dict) and isinstance(row.get("index"), int) for row in rows):
|
|
|
+ rows = sorted(rows, key=lambda row: int(row["index"]))
|
|
|
+ vectors = [row.get("embedding") if isinstance(row, dict) else None for row in rows]
|
|
|
+ if any(not isinstance(vector, list) or len(vector) != dimension for vector in vectors):
|
|
|
+ raise EmbeddingApiError("embedding response contains an invalid vector dimension")
|
|
|
+ matrix = np.asarray(vectors, dtype=np.float32)
|
|
|
+ if matrix.shape != (expected_rows, dimension) or not np.isfinite(matrix).all():
|
|
|
+ raise EmbeddingApiError("embedding response contains invalid numeric values")
|
|
|
+ usage = payload.get("usage") if isinstance(payload, dict) else None
|
|
|
+ tokens = usage.get("prompt_tokens", 0) if isinstance(usage, dict) else 0
|
|
|
+ if isinstance(tokens, bool) or not isinstance(tokens, int) or tokens < 0:
|
|
|
+ raise EmbeddingApiError("embedding response contains invalid token usage")
|
|
|
+ return matrix, tokens
|
|
|
+
|
|
|
+
|
|
|
+def _retry_after_seconds(response: httpx.Response) -> float | None:
|
|
|
+ value = response.headers.get("Retry-After")
|
|
|
+ if value is None:
|
|
|
+ return None
|
|
|
+ try:
|
|
|
+ return min(max(float(value), 0.0), 30.0)
|
|
|
+ except ValueError:
|
|
|
+ return None
|
|
|
+
|
|
|
+
|
|
|
+def _bounded_api_error(response: httpx.Response) -> str:
|
|
|
+ try:
|
|
|
+ payload = response.json()
|
|
|
+ if isinstance(payload, dict):
|
|
|
+ message = payload.get("message")
|
|
|
+ error = payload.get("error")
|
|
|
+ if not message and isinstance(error, dict):
|
|
|
+ message = error.get("message")
|
|
|
+ detail = str(message or "remote API rejected the request")[:300]
|
|
|
+ else:
|
|
|
+ detail = "remote API rejected the request"
|
|
|
+ except ValueError:
|
|
|
+ detail = "remote API returned a non-JSON error"
|
|
|
+ return f"embedding API returned HTTP {response.status_code}: {detail}"
|
|
|
+
|
|
|
+
|
|
|
+def _json_lines(rows: Sequence[Mapping[str, Any]]) -> bytes:
|
|
|
+ return b"".join(
|
|
|
+ (json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n").encode("utf-8")
|
|
|
+ for row in rows
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+def _directory_digest(root: Path) -> str:
|
|
|
+ digest = sha256()
|
|
|
+ for path in sorted(root.rglob("*.json")):
|
|
|
+ digest.update(path.relative_to(root).as_posix().encode("utf-8"))
|
|
|
+ digest.update(b"\0")
|
|
|
+ with path.open("rb") as handle:
|
|
|
+ for block in iter(lambda: handle.read(1024 * 1024), b""):
|
|
|
+ digest.update(block)
|
|
|
+ digest.update(b"\0")
|
|
|
+ return "sha256:" + digest.hexdigest()
|
|
|
+
|
|
|
+
|
|
|
+def _file_digest(path: Path) -> str:
|
|
|
+ digest = sha256()
|
|
|
+ with path.open("rb") as handle:
|
|
|
+ for block in iter(lambda: handle.read(1024 * 1024), b""):
|
|
|
+ digest.update(block)
|
|
|
+ return "sha256:" + digest.hexdigest()
|
|
|
+
|
|
|
+
|
|
|
+def _atomic_write_bytes(path: Path, value: bytes) -> None:
|
|
|
+ temporary = path.with_name(f".{path.name}.{uuid4().hex}.tmp")
|
|
|
+ temporary.write_bytes(value)
|
|
|
+ os.replace(temporary, path)
|
|
|
+
|
|
|
+
|
|
|
+def _atomic_write_json(path: Path, value: Mapping[str, Any]) -> None:
|
|
|
+ _atomic_write_bytes(
|
|
|
+ path,
|
|
|
+ (json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n").encode("utf-8"),
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+def _atomic_save_array(path: Path, value: Any) -> None:
|
|
|
+ temporary = path.with_name(f".{path.name}.{uuid4().hex}.tmp")
|
|
|
+ with temporary.open("wb") as handle:
|
|
|
+ np.save(handle, value)
|
|
|
+ os.replace(temporary, path)
|
|
|
+
|
|
|
+
|
|
|
+def _read_json_object(path: Path) -> dict[str, Any]:
|
|
|
+ try:
|
|
|
+ payload = json.loads(path.read_text(encoding="utf-8"))
|
|
|
+ except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
|
|
|
+ raise RuntimeError(f"build state is not valid UTF-8 JSON: {path}") from exc
|
|
|
+ if not isinstance(payload, dict):
|
|
|
+ raise RuntimeError(f"build state root must be an object: {path}")
|
|
|
+ return payload
|
|
|
+
|
|
|
+
|
|
|
+def _verify_vectors(path: Path, rows: int, dimension: int) -> tuple[float, float]:
|
|
|
+ matrix = np.load(path, mmap_mode="r")
|
|
|
+ if matrix.shape != (rows, dimension):
|
|
|
+ raise RuntimeError("published embedding matrix shape is invalid")
|
|
|
+ minimum = float("inf")
|
|
|
+ maximum = 0.0
|
|
|
+ for start in range(0, rows, 2_000):
|
|
|
+ values = matrix[start : start + 2_000].astype(np.float32)
|
|
|
+ if not np.isfinite(values).all():
|
|
|
+ raise RuntimeError("published embedding matrix contains a non-finite value")
|
|
|
+ norms = np.linalg.norm(values, axis=1)
|
|
|
+ if np.any(norms == 0):
|
|
|
+ raise RuntimeError("published embedding matrix contains a zero vector")
|
|
|
+ minimum = min(minimum, float(norms.min()))
|
|
|
+ maximum = max(maximum, float(norms.max()))
|
|
|
+ return round(minimum, 6), round(maximum, 6)
|
|
|
+
|
|
|
+
|
|
|
+def _publish_index(stage_root: Path, index_root: Path) -> None:
|
|
|
+ backup_root = index_root.parent / f".{index_root.name}-backup-{uuid4().hex}"
|
|
|
+ had_previous = index_root.exists()
|
|
|
+ if had_previous:
|
|
|
+ os.replace(index_root, backup_root)
|
|
|
+ try:
|
|
|
+ os.replace(stage_root, index_root)
|
|
|
+ except BaseException:
|
|
|
+ if had_previous and backup_root.exists() and not index_root.exists():
|
|
|
+ os.replace(backup_root, index_root)
|
|
|
+ raise
|
|
|
+ if backup_root.exists():
|
|
|
+ shutil.rmtree(backup_root)
|
|
|
+
|
|
|
+
|
|
|
+def _api_key(env_file: Path) -> str:
|
|
|
+ direct = os.environ.get("DASHSCOPE_API_KEY") or os.environ.get("SCRIPT_BUILD_EMBEDDING_API_KEY")
|
|
|
+ if direct:
|
|
|
+ return direct
|
|
|
+ values: dict[str, str] = {}
|
|
|
+ try:
|
|
|
+ lines = env_file.read_text(encoding="utf-8").splitlines()
|
|
|
+ except OSError as exc:
|
|
|
+ raise RuntimeError(f"embedding env file cannot be read: {env_file}") from exc
|
|
|
+ for raw_line in lines:
|
|
|
+ line = raw_line.strip()
|
|
|
+ if not line or line.startswith("#") or "=" not in line:
|
|
|
+ continue
|
|
|
+ key, configured_value = line.removeprefix("export ").split("=", 1)
|
|
|
+ values[key.strip()] = configured_value.strip().strip('"').strip("'")
|
|
|
+ api_key = values.get("DASHSCOPE_API_KEY") or values.get("SCRIPT_BUILD_EMBEDDING_API_KEY")
|
|
|
+ if not api_key:
|
|
|
+ raise RuntimeError("DASHSCOPE_API_KEY or SCRIPT_BUILD_EMBEDDING_API_KEY is required")
|
|
|
+ return api_key
|
|
|
+
|
|
|
+
|
|
|
+def main() -> None:
|
|
|
+ parser = argparse.ArgumentParser(description="全量重建本项目的脚本解构向量索引")
|
|
|
+ parser.add_argument("--raw-root", type=Path, required=True)
|
|
|
+ parser.add_argument("--index-root", type=Path, required=True)
|
|
|
+ parser.add_argument("--env-file", type=Path, default=Path(".env"))
|
|
|
+ parser.add_argument(
|
|
|
+ "--endpoint",
|
|
|
+ default="https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings",
|
|
|
+ )
|
|
|
+ parser.add_argument("--model", default="text-embedding-v4")
|
|
|
+ parser.add_argument("--dimension", type=int, default=512)
|
|
|
+ parser.add_argument("--batch-size", type=int, default=10)
|
|
|
+ parser.add_argument("--concurrency", type=int, default=8)
|
|
|
+ parser.add_argument("--shard-size", type=int, default=1_000)
|
|
|
+ parser.add_argument("--storage-dtype", choices=("float16", "float32"), default="float16")
|
|
|
+ parser.add_argument("--reset-work", action="store_true")
|
|
|
+ args = parser.parse_args()
|
|
|
+
|
|
|
+ async def run() -> DecodeIndexBuildResult:
|
|
|
+ timeout = httpx.Timeout(60.0, connect=15.0)
|
|
|
+ async with httpx.AsyncClient(timeout=timeout, trust_env=False) as http:
|
|
|
+ client = DashScopeEmbeddingClient(
|
|
|
+ http,
|
|
|
+ api_key=_api_key(args.env_file),
|
|
|
+ endpoint=args.endpoint,
|
|
|
+ model=args.model,
|
|
|
+ dimension=args.dimension,
|
|
|
+ concurrency=args.concurrency,
|
|
|
+ )
|
|
|
+ return await rebuild_decode_index(
|
|
|
+ raw_root=args.raw_root,
|
|
|
+ index_root=args.index_root,
|
|
|
+ client=client,
|
|
|
+ batch_size=args.batch_size,
|
|
|
+ shard_size=args.shard_size,
|
|
|
+ storage_dtype=args.storage_dtype,
|
|
|
+ reset_work=args.reset_work,
|
|
|
+ )
|
|
|
+
|
|
|
+ result = asyncio.run(run())
|
|
|
+ print(json.dumps(result.manifest, ensure_ascii=False, indent=2, sort_keys=True))
|
|
|
+ print(f"✅ 已原子发布索引: {result.index_root}")
|
|
|
+
|
|
|
+
|
|
|
+if __name__ == "__main__":
|
|
|
+ main()
|
SamLee