Repository navigation
Expand file tree
/
Copy pathembedder.py
More file actions
329 lines (275 loc) · 11.9 KB
/
Copy pathembedder.py
File metadata and controls
329 lines (275 loc) · 11.9 KB
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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
from __future__ import annotations
import hashlib
import json
import os
from functools import lru_cache
from pathlib import Path
import httpx
from loguru import logger
from . import constants as cs
from . import exceptions as ex
from . import logs as ls
from .config import API_KEY_INFO, settings
from .utils.dependencies import has_torch, has_transformers
def _cache_namespace() -> str:
# Embeddings from different providers/models live in different vector
# spaces; namespacing cache keys prevents a provider or model switch
# from replaying stale vectors of the wrong space.
if settings.EMBEDDING_PROVIDER == cs.EmbeddingProvider.OPENAI:
namespace = f"{cs.EmbeddingProvider.OPENAI}:{settings.OPENAI_EMBEDDING_MODEL}"
if settings.OPENAI_EMBEDDING_DIMENSIONS is not None:
namespace = f"{namespace}:{settings.OPENAI_EMBEDDING_DIMENSIONS}"
return namespace
return f"{cs.EmbeddingProvider.UNIXCODER}:{cs.UNIXCODER_MODEL}"
class EmbeddingCache:
__slots__ = ("_cache", "_path")
def __init__(self, path: Path | None = None) -> None:
self._cache: dict[str, list[float]] = {}
self._path = path
@staticmethod
def _content_hash(content: str) -> str:
return hashlib.sha256(f"{_cache_namespace()}\x00{content}".encode()).hexdigest()
def get(self, content: str) -> list[float] | None:
return self._cache.get(self._content_hash(content))
def put(self, content: str, embedding: list[float]) -> None:
self._cache[self._content_hash(content)] = embedding
def get_many(self, snippets: list[str]) -> dict[int, list[float]]:
results: dict[int, list[float]] = {}
for i, snippet in enumerate(snippets):
if (cached := self.get(snippet)) is not None:
results[i] = cached
return results
def put_many(self, snippets: list[str], embeddings: list[list[float]]) -> None:
for snippet, embedding in zip(snippets, embeddings):
self.put(snippet, embedding)
def save(self) -> None:
if self._path is None:
return
try:
self._path.parent.mkdir(parents=True, exist_ok=True)
with self._path.open("w", encoding="utf-8") as f:
json.dump(self._cache, f)
except Exception as e:
logger.warning(ls.EMBEDDING_CACHE_SAVE_FAILED, path=self._path, error=e)
def load(self) -> None:
if self._path is None or not self._path.exists():
return
try:
with self._path.open("r", encoding="utf-8") as f:
self._cache = json.load(f)
logger.debug(
ls.EMBEDDING_CACHE_LOADED, count=len(self._cache), path=self._path
)
except Exception as e:
logger.warning(ls.EMBEDDING_CACHE_LOAD_FAILED, path=self._path, error=e)
self._cache = {}
def clear(self) -> None:
self._cache.clear()
def __len__(self) -> int:
return len(self._cache)
_embedding_cache: EmbeddingCache | None = None
def get_embedding_cache() -> EmbeddingCache:
global _embedding_cache
if _embedding_cache is None:
cache_path = Path(settings.QDRANT_DB_PATH) / cs.EMBEDDING_CACHE_FILENAME
_embedding_cache = EmbeddingCache(path=cache_path)
_embedding_cache.load()
return _embedding_cache
def clear_embedding_cache() -> None:
global _embedding_cache
if _embedding_cache is not None:
_embedding_cache.clear()
_embedding_cache = None
def _openai_client() -> httpx.Client:
headers: dict[str, str] = {}
api_key = settings.OPENAI_EMBEDDING_API_KEY or os.environ.get(
API_KEY_INFO[cs.Provider.OPENAI]["env_var"]
)
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
return httpx.Client(
base_url=settings.OPENAI_EMBEDDING_BASE_URL,
headers=headers,
timeout=settings.OPENAI_EMBEDDING_TIMEOUT,
)
def _openai_embed_batch(snippets: list[str]) -> list[list[float]]:
embeddings: list[list[float]] = []
batch_size = settings.OPENAI_EMBEDDING_BATCH_SIZE
with _openai_client() as client:
for start in range(0, len(snippets), batch_size):
batch = snippets[start : start + batch_size]
payload: dict[str, object] = {
"model": settings.OPENAI_EMBEDDING_MODEL,
"input": batch,
}
if settings.OPENAI_EMBEDDING_DIMENSIONS is not None:
payload["dimensions"] = settings.OPENAI_EMBEDDING_DIMENSIONS
response = client.post(cs.OPENAI_EMBEDDINGS_PATH, json=payload)
if response.status_code != cs.HTTP_OK:
raise RuntimeError(
ex.OPENAI_EMBEDDING_HTTP_ERROR.format(
status=response.status_code, body=response.text[:500]
)
)
embeddings.extend(_parse_embedding_rows(response, len(batch)))
return embeddings
def _parse_embedding_rows(response: httpx.Response, expected: int) -> list[list[float]]:
try:
rows = response.json()["data"]
except (json.JSONDecodeError, KeyError, TypeError) as e:
raise RuntimeError(
ex.OPENAI_EMBEDDING_MALFORMED_RESPONSE.format(error=e)
) from e
if not isinstance(rows, list):
raise RuntimeError(
ex.OPENAI_EMBEDDING_MALFORMED_RESPONSE.format(error="'data' is not a list")
)
if len(rows) != expected:
raise RuntimeError(
ex.OPENAI_EMBEDDING_COUNT_MISMATCH.format(got=len(rows), expected=expected)
)
# Place each row at its declared index instead of sorting: a buggy
# server emitting duplicate or out-of-range indices must fail loudly
# rather than pair snippets with the wrong vectors.
placed: list[list[float] | None] = [None] * expected
for row in rows:
try:
index, embedding = row["index"], row["embedding"]
except (KeyError, TypeError) as e:
raise RuntimeError(
ex.OPENAI_EMBEDDING_MALFORMED_RESPONSE.format(error=e)
) from e
if (
not isinstance(index, int)
or not 0 <= index < expected
or placed[index] is not None
):
raise RuntimeError(ex.OPENAI_EMBEDDING_BAD_INDEX.format(index=index))
placed[index] = embedding
return [row for row in placed if row is not None]
if has_torch() and has_transformers():
import numpy as np
import torch
from numpy.typing import NDArray
from .unixcoder import UniXcoder
def _device_available(device: cs.EmbeddingDevice) -> bool:
match device:
case cs.EmbeddingDevice.CUDA:
return torch.cuda.is_available()
case cs.EmbeddingDevice.MPS:
return torch.backends.mps.is_available()
case _:
return True
def _select_device() -> cs.EmbeddingDevice:
if (override := settings.EMBEDDING_DEVICE) is not None:
if _device_available(override):
return override
logger.warning(ls.EMBEDDING_DEVICE_UNAVAILABLE.format(device=override))
if torch.cuda.is_available():
return cs.EmbeddingDevice.CUDA
if torch.backends.mps.is_available():
return cs.EmbeddingDevice.MPS
return cs.EmbeddingDevice.CPU
_batches_since_cache_drop = 0
def _sync_after_batch(device: torch.device | str) -> None:
# MPS wedges inside Metal's waitUntilCompleted when command buffers
# accumulate across thousands of batches (issue #689); draining the
# stream after every batch keeps each command buffer short-lived,
# and a periodic (not per-batch: ~21% throughput cost) cache drop
# bounds Metal allocator growth over monorepo-scale runs.
if torch.device(device).type != cs.EmbeddingDevice.MPS:
return
torch.mps.synchronize()
global _batches_since_cache_drop
_batches_since_cache_drop += 1
if _batches_since_cache_drop >= cs.EMBEDDING_MPS_CACHE_DROP_INTERVAL:
torch.mps.empty_cache()
_batches_since_cache_drop = 0
@lru_cache(maxsize=1)
def get_model() -> UniXcoder:
model = UniXcoder(cs.UNIXCODER_MODEL)
model.eval()
device = _select_device()
if device != cs.EmbeddingDevice.CPU:
model = model.to(device)
return model
def _unixcoder_embed_code(code: str, max_length: int | None) -> list[float]:
if max_length is None:
max_length = settings.EMBEDDING_MAX_LENGTH
model = get_model()
device = next(model.parameters()).device
tokens = model.tokenize([code], max_length=max_length)
tokens_tensor = torch.tensor(tokens).to(device)
with torch.no_grad():
_, sentence_embeddings = model(tokens_tensor)
embedding: NDArray[np.float32] = sentence_embeddings.cpu().numpy()
_sync_after_batch(device)
result: list[float] = embedding[0].tolist()
return result
def _unixcoder_embed_batch(
snippets: list[str], max_length: int | None, batch_size: int
) -> list[list[float]]:
if max_length is None:
max_length = settings.EMBEDDING_MAX_LENGTH
model = get_model()
device = next(model.parameters()).device
all_new_embeddings: list[list[float]] = []
for start in range(0, len(snippets), batch_size):
batch = snippets[start : start + batch_size]
tokens_list = model.tokenize(batch, max_length=max_length, padding=True)
tokens_tensor = torch.tensor(tokens_list).to(device)
with torch.no_grad():
_, sentence_embeddings = model(tokens_tensor)
batch_np: NDArray[np.float32] = sentence_embeddings.cpu().numpy()
_sync_after_batch(device)
for row in batch_np:
all_new_embeddings.append(row.tolist())
return all_new_embeddings
else:
def _unixcoder_embed_code(code: str, max_length: int | None) -> list[float]:
raise RuntimeError(ex.SEMANTIC_EXTRA)
def _unixcoder_embed_batch(
snippets: list[str], max_length: int | None, batch_size: int
) -> list[list[float]]:
raise RuntimeError(ex.SEMANTIC_EXTRA)
def embed_code(code: str, max_length: int | None = None) -> list[float]:
cache = get_embedding_cache()
if (cached := cache.get(code)) is not None:
return cached
try:
if settings.EMBEDDING_PROVIDER == cs.EmbeddingProvider.OPENAI:
result = _openai_embed_batch([code])[0]
else:
result = _unixcoder_embed_code(code, max_length)
except Exception:
logger.exception(ls.EMBEDDING_SNIPPET_FAILED, length=len(code))
raise
cache.put(code, result)
return result
def embed_code_batch(
snippets: list[str],
max_length: int | None = None,
batch_size: int = cs.EMBEDDING_DEFAULT_BATCH_SIZE,
) -> list[list[float]]:
if not snippets:
return []
cache = get_embedding_cache()
cached_results = cache.get_many(snippets)
if len(cached_results) == len(snippets):
logger.debug(ls.EMBEDDING_CACHE_HIT, count=len(snippets))
return [cached_results[i] for i in range(len(snippets))]
uncached_indices = [i for i in range(len(snippets)) if i not in cached_results]
uncached_snippets = [snippets[i] for i in uncached_indices]
if settings.EMBEDDING_PROVIDER == cs.EmbeddingProvider.OPENAI:
all_new_embeddings = _openai_embed_batch(uncached_snippets)
else:
all_new_embeddings = _unixcoder_embed_batch(
uncached_snippets, max_length, batch_size
)
cache.put_many(uncached_snippets, all_new_embeddings)
results: list[list[float]] = [[] for _ in snippets]
for i, emb in cached_results.items():
results[i] = emb
for idx, orig_i in enumerate(uncached_indices):
results[orig_i] = all_new_embeddings[idx]
return results