Add missing import in NvidiaGenerator docstring (#559)

* Add missing import in NvidiaGenerator docstring * Export NvidiaGeneratorModel in nvidia package
deepset-ai · Mar 7, 2024 · 6e00364 · 6e00364
1 parent ead381e
commit 6e00364
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/__init__.py b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/__init__.py
@@ -2,5 +2,6 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 from .generator import NvidiaGenerator
+from .models import NvidiaGeneratorModel
 
-__all__ = ["NvidiaGenerator"]
+__all__ = ["NvidiaGenerator", "NvidiaGeneratorModel"]
diff --git a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py
@@ -19,7 +19,7 @@ class NvidiaGenerator:
 
     Usage example:
     ```python
-    from haystack_integrations.components.generators.nvidia import NvidiaGenerator
+    from haystack_integrations.components.generators.nvidia import NvidiaGenerator, NvidiaGeneratorModel
 
     generator = NvidiaGenerator(
         model=NvidiaGeneratorModel.NV_LLAMA2_RLHF_70B,