kgmodule-utils 0.4.7__tar.gz → 0.4.8__tar.gz
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.
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/PKG-INFO +2 -2
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/README.md +1 -1
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/pyproject.toml +1 -1
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/__init__.py +1 -1
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/corpus_embedder.py +58 -3
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/LICENSE +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/embed.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/embedder.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/extractor.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/module.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/pipeline.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/py.typed +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/retrieval/__init__.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/retrieval/hits.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/semantic.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/snapshots/__init__.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/snapshots/manager.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/snapshots/models.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/specs.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/store.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/synthesis/__init__.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/synthesis/_config.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/synthesis/_image.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/synthesis/_text.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/synthesis/factory.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/worker/__init__.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/worker/client.py +0 -0
- {kgmodule_utils-0.4.7 → kgmodule_utils-0.4.8}/src/kg_utils/worker/ops.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: kgmodule-utils
|
|
3
|
-
Version: 0.4.
|
|
3
|
+
Version: 0.4.8
|
|
4
4
|
Summary: Shared types, graph store, semantic index, and pipeline base for the KGModule SDK
|
|
5
5
|
License: Elastic-2.0
|
|
6
6
|
License-File: LICENSE
|
|
@@ -36,7 +36,7 @@ Description-Content-Type: text/markdown
|
|
|
36
36
|
|
|
37
37
|
[](https://www.python.org/)
|
|
38
38
|
[](https://www.elastic.co/licensing/elastic-license)
|
|
39
|
-
[](https://github.com/Flux-Frontiers/KG_utils/releases)
|
|
40
40
|
[](https://github.com/Flux-Frontiers/KG_utils/actions/workflows/ci.yml)
|
|
41
41
|
[](https://python-poetry.org/)
|
|
42
42
|
|
|
@@ -1,7 +1,7 @@
|
|
|
1
1
|
|
|
2
2
|
[](https://www.python.org/)
|
|
3
3
|
[](https://www.elastic.co/licensing/elastic-license)
|
|
4
|
-
[](https://github.com/Flux-Frontiers/KG_utils/releases)
|
|
5
5
|
[](https://github.com/Flux-Frontiers/KG_utils/actions/workflows/ci.yml)
|
|
6
6
|
[](https://python-poetry.org/)
|
|
7
7
|
|
|
@@ -10,7 +10,7 @@ build-backend = "poetry.core.masonry.api"
|
|
|
10
10
|
|
|
11
11
|
[project]
|
|
12
12
|
name = "kgmodule-utils"
|
|
13
|
-
version = "0.4.
|
|
13
|
+
version = "0.4.8"
|
|
14
14
|
description = "Shared types, graph store, semantic index, and pipeline base for the KGModule SDK"
|
|
15
15
|
readme = "README.md"
|
|
16
16
|
license = { text = "Elastic-2.0" }
|
|
@@ -123,6 +123,26 @@ def _embed_shard(args: tuple) -> tuple[int, list[list[float]]]:
|
|
|
123
123
|
return worker_id, [np.asarray(v, dtype="float32").tolist() for v in all_vecs]
|
|
124
124
|
|
|
125
125
|
|
|
126
|
+
class _InlineProgress:
|
|
127
|
+
"""In-process stand-in for the multiprocessing progress queue.
|
|
128
|
+
|
|
129
|
+
Lets :func:`_embed_shard` report per-batch progress when it runs in the
|
|
130
|
+
main process (sequential/GPU path) — same ``put`` protocol, no queue.
|
|
131
|
+
|
|
132
|
+
:param progress: Live ``rich.progress.Progress`` instance.
|
|
133
|
+
:param task_id: Task to advance.
|
|
134
|
+
"""
|
|
135
|
+
|
|
136
|
+
def __init__(self, progress, task_id) -> None:
|
|
137
|
+
self._progress = progress
|
|
138
|
+
self._task_id = task_id
|
|
139
|
+
|
|
140
|
+
def put(self, item: int | None) -> None:
|
|
141
|
+
"""Advance the bar by *item* texts; ignore the end-of-shard sentinel."""
|
|
142
|
+
if item is not None:
|
|
143
|
+
self._progress.advance(self._task_id, item)
|
|
144
|
+
|
|
145
|
+
|
|
126
146
|
# ============================================================================
|
|
127
147
|
# Embedding cache
|
|
128
148
|
# ============================================================================
|
|
@@ -167,7 +187,11 @@ class CorpusEmbedder:
|
|
|
167
187
|
"""Multi-process corpus embedding engine.
|
|
168
188
|
|
|
169
189
|
:param model_name: HuggingFace model name.
|
|
170
|
-
:param n_workers: Number of parallel workers (default:
|
|
190
|
+
:param n_workers: Number of parallel workers (default: ``min(4, cpu_count // 2)``).
|
|
191
|
+
Each CPU worker loads its own full model copy plus a torch runtime
|
|
192
|
+
(~1.2 GB for bge-small, more for mpnet-class models), so raise this
|
|
193
|
+
past 4 only when memory allows — throughput is I/O + accumulator
|
|
194
|
+
bound well before then.
|
|
171
195
|
:param batch_size: Per-worker batch size.
|
|
172
196
|
:param device: Embedding device (``"cpu"``/``"mps"``/``"cuda"``). ``None``
|
|
173
197
|
resolves via ``KG_EMBED_DEVICE`` then auto-detect. A GPU device forces
|
|
@@ -185,7 +209,7 @@ class CorpusEmbedder:
|
|
|
185
209
|
) -> None:
|
|
186
210
|
"""Configure the embedding engine; workers are spawned lazily during :meth:`embed`."""
|
|
187
211
|
self.model_name = model_name
|
|
188
|
-
self.n_workers = n_workers or max(1, (os.cpu_count() or 2) // 2)
|
|
212
|
+
self.n_workers = n_workers or min(4, max(1, (os.cpu_count() or 2) // 2))
|
|
189
213
|
self.batch_size = batch_size
|
|
190
214
|
self.device = resolve_device(device)
|
|
191
215
|
|
|
@@ -250,7 +274,38 @@ class CorpusEmbedder:
|
|
|
250
274
|
|
|
251
275
|
def _embed_sequential(self, texts: list[str]) -> list[list[float]]:
|
|
252
276
|
"""Embed in the main process (small inputs, single worker, or GPU device)."""
|
|
253
|
-
|
|
277
|
+
from rich.progress import ( # pylint: disable=import-outside-toplevel
|
|
278
|
+
BarColumn,
|
|
279
|
+
MofNCompleteColumn,
|
|
280
|
+
Progress,
|
|
281
|
+
SpinnerColumn,
|
|
282
|
+
TextColumn,
|
|
283
|
+
TimeElapsedColumn,
|
|
284
|
+
TimeRemainingColumn,
|
|
285
|
+
)
|
|
286
|
+
|
|
287
|
+
with Progress(
|
|
288
|
+
SpinnerColumn(),
|
|
289
|
+
TextColumn("[progress.description]{task.description}"),
|
|
290
|
+
BarColumn(),
|
|
291
|
+
MofNCompleteColumn(),
|
|
292
|
+
TimeElapsedColumn(),
|
|
293
|
+
TimeRemainingColumn(),
|
|
294
|
+
) as progress:
|
|
295
|
+
task = progress.add_task(
|
|
296
|
+
f" Embedding (single process, {self.device or 'auto'})",
|
|
297
|
+
total=len(texts),
|
|
298
|
+
)
|
|
299
|
+
_, vectors = _embed_shard(
|
|
300
|
+
(
|
|
301
|
+
texts,
|
|
302
|
+
self.model_name,
|
|
303
|
+
self.batch_size,
|
|
304
|
+
0,
|
|
305
|
+
_InlineProgress(progress, task),
|
|
306
|
+
self.device,
|
|
307
|
+
)
|
|
308
|
+
)
|
|
254
309
|
return vectors
|
|
255
310
|
|
|
256
311
|
def _embed_parallel(self, texts: list[str]) -> list[list[float]]:
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|