File size: 3,265 Bytes
ee7e79e e1dced2 ee7e79e e1dced2 ee7e79e e1dced2 ee7e79e e1dced2 ee7e79e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | """Client for the university's OpenAI-compatible Qwen3-Embedding-4B endpoint.
The endpoint runs a 4B model, so a single request carrying an entire weekly
delta (or a first-run full anthology) easily exceeds the read timeout. We
therefore chunk the input and retry the transient timeouts the openai SDK
does not retry on its own (it retries 429/5xx but treats read timeouts as
non-retryable connection errors).
"""
import logging
import time
import numpy as np
from openai import APIConnectionError, OpenAI, OpenAIError
from tqdm import tqdm
logger = logging.getLogger(__name__)
# Max texts per embedding request. Small enough that each call to the 4B
# model finishes well inside the read timeout, even on a slow week.
_BATCH_SIZE = 32
# Number of attempts per chunk before giving up (and propagating before any
# writes in run_sync, per its "fail before writes" contract).
_MAX_RETRIES = 3
_RETRY_BACKOFF_SECONDS = 2.0
class EmbeddingsClient:
def __init__(self, base_url: str, api_key: str, model: str = "qwen3-embedding-4b",
timeout: float = 60.0):
# The hosted qwen3-embedding-4b does NOT support matryoshka representation,
# so we must not send `dimensions` — sending it 400s. The model returns its
# fixed native dimensionality; the index dim is derived from the returned
# vectors in run_sync rather than assumed here.
self.model = model
self._client = OpenAI(base_url=base_url, api_key=api_key, timeout=timeout)
def embed_batch(self, texts: list[str]) -> list[np.ndarray]:
if not texts:
return []
total_chunks = (len(texts) + _BATCH_SIZE - 1) // _BATCH_SIZE
vectors: list[np.ndarray] = []
starts = range(0, len(texts), _BATCH_SIZE)
for start in tqdm(starts, total=total_chunks, desc="embedding", unit="batch"):
chunk = texts[start:start + _BATCH_SIZE]
vectors.extend(self._embed_chunk(chunk))
return vectors
def _embed_chunk(self, chunk: list[str]) -> list[np.ndarray]:
last_error: Exception | None = None
for attempt in range(1, _MAX_RETRIES + 1):
try:
response = self._client.embeddings.create(
model=self.model, input=chunk
)
except APIConnectionError as e:
# Read timeouts (APITimeoutError is a subclass) and other
# transient connection errors: the SDK won't retry, so we do.
last_error = e
if attempt == _MAX_RETRIES:
break
logger.warning(
"embedding request failed (attempt %d/%d), retrying in %.1fs: %s",
attempt, _MAX_RETRIES, _RETRY_BACKOFF_SECONDS, e,
)
time.sleep(_RETRY_BACKOFF_SECONDS)
continue
except OpenAIError as e:
raise RuntimeError(f"Embedding API request failed: {e}") from e
data = sorted(response.data, key=lambda d: d.index)
return [np.array(d.embedding, dtype=np.float32) for d in data]
raise RuntimeError(
f"Embedding API request failed after {_MAX_RETRIES} attempts: {last_error}"
) from last_error
|