Skip to content

feature(embeddings): add embedding workflow + cpu/gpu workers - #123

Merged
winsomeglint merged 2 commits into
mainfrom
lsummerbll/extract-worker-embedding
Oct 21, 2026
Merged

winsomeglint merged 2 commits into
mainfrom
lsummerbll/extract-worker-embedding

Conversation

@winsomeglint

Copy link
Copy Markdown
Contributor
  • adds an HFEmbedder working with HuggingFace models exposing a jina-embeddings-like encode method (task + prompt name)
  • adds publish jobs for datashare-extract-cpu-worker and datashare-extract-gpu-worker

@winsomeglint
winsomeglint force-pushed the lsummerbll/extract-worker-embedding branch from 79c3b4f to f0ccb3c Compare October 21, 2026 14:37
Comment thread workers/extract-worker/extract_worker/config.py
output_path: Path,
progress: AsyncProgressRateHandler | None = None,
) -> EmbeddingResponse:
docs = list(read_jsonl_as(batch, PROCESSED_FILE_TA))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To avoid heartbeat timeout we can do something like using the async API:
[d async for d in async_read_jsonl_as()]

if await es_client.indices.exists(index=index):
return
logger.info("creating embeddings index %s...", index)
await es_client.indices.create(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think we'll have permission to do so, maybe we can just raise
for tests we can create the index in the fixtures

# run chunking and inference in a separate thread to avoid
# blocking the event loop
chunks = await asyncio.to_thread(
_embed_doc,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if heartbeat timeouts are recurrent we can also make _embed_doc an async corouting, use async_open when reading files and isolate calls to text embedding with asyncio.to_thread

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe it would be better to put heartbeating in a background task?

@winsomeglint
winsomeglint merged commit aad703e into main Oct 21, 2026
3 checks passed
@winsomeglint
winsomeglint deleted the lsummerbll/extract-worker-embedding branch October 21, 2026 16:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants