feat: add device config (cpu/cuda) for all enrichment models#1
Open
tmaiaroto wants to merge 6 commits into
Open
feat: add device config (cpu/cuda) for all enrichment models#1tmaiaroto wants to merge 6 commits into
tmaiaroto wants to merge 6 commits into
Conversation
- Add 'device' field to EmbeddingConfig, SummarizationConfig, and ClassifierConfig - Update EnrichmentConfig.from_dict() to pass through device configs - Update EnrichmentService to use device config when loading models - Update example YAML to document the new device option - All models now respect the configured device (cpu or cuda) instead of hardcoded cpu
- Add 'mps' as valid device option for Apple Silicon GPUs - Add 'auto' option to auto-detect best available device (cuda > mps > cpu) - Update example YAML to use 'auto' by default - Falls back to cpu if requested device isn't available
- Test default device is cpu - Test device from config - Test kwargs override config - Test invalid device falls back to cpu - Test auto-detect for cuda, mps, and cpu fallback
- Add Windows-specific installation instructions to README - Use %LOCALAPPDATA% path placeholder (no hardcoded username) - Add test_device_config.py with 7 tests for cpu/cuda/mps/auto - Coverage: enrichment.py 70%, schemas.py 78%, total 72%
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds configurable device selection (cpu/cuda) for all enrichment models, allowing users to offload inference to GPU if available.
Changes
"cpu")EnrichmentConfig.from_dict()to pass through device configsEnrichmentServiceto use device config when loading modelsdeviceoption for each modelConfig Example
Notes
"cpu"for backward compatibility"cpu"and"cuda"are valid values (invalid values fall back to cpu with a warning)deviceconfig option inretrieval.rewriter