Embeddings: Update config, args, and parameter names

Use embeddings_device as the parameter for device to remove ambiguity.

Signed-off-by: kingbri <bdashore3@proton.me>
This commit is contained in:
kingbri
2024-07-30 15:32:26 -04:00
parent bfa011e0ce
commit dc3dcc9c0d
5 changed files with 43 additions and 9 deletions

View File

@@ -35,7 +35,7 @@ class InfinityContainer:
self.model_is_loading = True
# Use cpu by default
device = unwrap(kwargs.get("device"), "cpu")
device = unwrap(kwargs.get("embeddings_device"), "cpu")
engine_args = EngineArgs(
model_name_or_path=str(self.model_dir),