Skip to content

feat: add device config (cpu/cuda) for all enrichment models#1

Open
tmaiaroto wants to merge 6 commits into
mainfrom
feature/model-device-config
Open

feat: add device config (cpu/cuda) for all enrichment models#1
tmaiaroto wants to merge 6 commits into
mainfrom
feature/model-device-config

Conversation

@tmaiaroto
Copy link
Copy Markdown
Contributor

Summary

Adds configurable device selection (cpu/cuda) for all enrichment models, allowing users to offload inference to GPU if available.

Changes

  • schemas.py: Added field to , , and (default: "cpu")
  • enrichment.py:
    • Updated EnrichmentConfig.from_dict() to pass through device configs
    • Updated EnrichmentService to use device config when loading models
    • Added validation to reject invalid device values (falls back to cpu)
  • reveriecore.yaml.example: Documented the new device option for each model

Config Example

enrichment:
  classifier:
    model: "MoritzLaurer/mDeBERTa-v3-base-mnli-xnli"
    device: "cuda"  # or "cpu"
  
  embedding:
    model: "BAAI/bge-small-en-v1.5"
    device: "cuda"  # or "cpu"
  
  summarization:
    model: "sshleifer/distilbart-cnn-12-6"
    device: "cuda"  # or "cpu"

Notes

  • Defaults to "cpu" for backward compatibility
  • Only "cpu" and "cuda" are valid values (invalid values fall back to cpu with a warning)
  • The query rewriter (GGUF model) already had a device config option in retrieval.rewriter

tmaiaroto added 6 commits May 29, 2026 22:43
- 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%
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.

1 participant