Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

import os
import warnings
from dataclasses import replace
from typing import Any, Optional, Union

from haystack import Document, component, default_from_dict, default_to_dict, logging
Expand Down Expand Up @@ -49,7 +50,7 @@ def __init__(
embedding_separator: str = "\n",
truncate: Optional[Union[EmbeddingTruncateMode, str]] = None,
timeout: Optional[float] = None,
):
) -> None:
"""
Create a NvidiaTextEmbedder component.

Expand Down Expand Up @@ -108,7 +109,7 @@ def __init__(
def class_name(cls) -> str:
return "NvidiaDocumentEmbedder"

def default_model(self):
def default_model(self) -> None:
"""Set default model in local NIM mode."""
valid_models = [
model.id for model in self.available_models if not model.base_model or model.base_model == model.id
Expand All @@ -129,7 +130,7 @@ def default_model(self):
error_message = "No locally hosted model was found."
raise ValueError(error_message)

def warm_up(self):
def warm_up(self) -> None:
"""
Initializes the component.
"""
Expand Down Expand Up @@ -246,14 +247,12 @@ def run(self, documents: list[Document]) -> dict[str, Union[list[Document], dict
A dictionary with the following keys and values:
- `documents` - List of processed Documents with embeddings.
- `meta` - Metadata on usage statistics, etc.
:raises RuntimeError:
If the component was not initialized.
:raises TypeError:
If the input is not a string.
If the input is not a list of Documents.
"""
if not self._initialized:
msg = "The embedding model has not been loaded. Please call warm_up() before running."
raise RuntimeError(msg)
self.warm_up()

elif not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
msg = (
"NvidiaDocumentEmbedder expects a list of Documents as input."
Expand All @@ -267,7 +266,9 @@ def run(self, documents: list[Document]) -> dict[str, Union[list[Document], dict

texts_to_embed = self._prepare_texts_to_embed(documents)
embeddings, metadata = self._embed_batch(texts_to_embed, self.batch_size)

new_documents = []
for doc, emb in zip(documents, embeddings):
doc.embedding = emb
new_documents.append(replace(doc, embedding=emb))

return {"documents": documents, "meta": metadata}
return {"documents": new_documents, "meta": metadata}
Original file line number Diff line number Diff line change
Expand Up @@ -197,14 +197,14 @@ def run(self, text: str) -> dict[str, Union[list[float], dict[str, Any]]]:
A dictionary with the following keys and values:
- `embedding` - Embedding of the text.
- `meta` - Metadata on usage statistics, etc.
:raises RuntimeError:
If the component was not initialized.
:raises TypeError:
If the input is not a string.
:raises ValueError:
If the input string is empty.
"""
if not self._initialized:
msg = "The embedding model has not been loaded. Please call warm_up() before running."
raise RuntimeError(msg)
self.warm_up()

elif not isinstance(text, str):
msg = (
"NvidiaTextEmbedder expects a string as an input."
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ def __init__(
timeout: Optional[float] = None,
max_retries: Optional[int] = None,
http_client_kwargs: Optional[dict[str, Any]] = None,
):
) -> None:
"""
Creates an instance of NvidiaChatGenerator.

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ def __init__(
api_key: Optional[Secret] = Secret.from_env_var("NVIDIA_API_KEY"),
model_arguments: Optional[dict[str, Any]] = None,
timeout: Optional[float] = None,
):
) -> None:
"""
Create a NvidiaGenerator component.

Expand Down Expand Up @@ -90,7 +90,7 @@ def __init__(
def class_name(cls) -> str:
return "NvidiaGenerator"

def default_model(self):
def default_model(self) -> None:
"""Set default model in local NIM mode."""
valid_models = [
model.id for model in self.available_models if not model.base_model or model.base_model == model.id
Expand All @@ -111,7 +111,7 @@ def default_model(self):
error_message = "No locally hosted model was found."
raise ValueError(error_message)

def warm_up(self):
def warm_up(self) -> None:
"""
Initializes the component.
"""
Expand Down Expand Up @@ -183,8 +183,7 @@ def run(self, prompt: str) -> dict[str, Union[list[str], list[dict[str, Any]]]]:
- `meta` - Metadata for each reply.
"""
if self.backend is None:
msg = "The generation model has not been loaded. Call warm_up() before running."
raise RuntimeError(msg)
self.warm_up()

assert self.backend is not None
replies, meta = self.backend.generate(prompt=prompt)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ def __init__(
meta_fields_to_embed: Optional[list[str]] = None,
embedding_separator: str = "\n",
timeout: Optional[float] = None,
):
) -> None:
"""
Create a NvidiaRanker component.

Expand Down Expand Up @@ -155,7 +155,7 @@ def from_dict(cls, data: dict[str, Any]) -> "NvidiaRanker":
deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"])
return default_from_dict(cls, data)

def warm_up(self):
def warm_up(self) -> None:
"""
Initialize the ranker.

Expand Down Expand Up @@ -192,27 +192,21 @@ def _prepare_documents_to_embed(self, documents: list[Document]) -> list[str]:
return document_texts

@component.output_types(documents=list[Document])
def run(
self,
query: str,
documents: list[Document],
top_k: Optional[int] = None,
) -> dict[str, list[Document]]:
def run(self, query: str, documents: list[Document], top_k: Optional[int] = None) -> dict[str, list[Document]]:
"""
Rank a list of documents based on a given query.

:param query: The query to rank the documents against.
:param documents: The list of documents to rank.
:param top_k: The number of documents to return.

:raises RuntimeError: If the ranker has not been loaded.
:raises TypeError: If the arguments are of the wrong type.

:returns: A dictionary containing the ranked documents.
"""
if not self._initialized:
msg = "The ranker has not been loaded. Please call warm_up() before running."
raise RuntimeError(msg)
self.warm_up()

if not isinstance(query, str):
msg = "NvidiaRanker expects the `query` parameter to be a string."
raise TypeError(msg)
Expand Down
Original file line number Diff line number Diff line change
@@ -1,3 +1,7 @@
# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

from enum import Enum


Expand Down
Original file line number Diff line number Diff line change
@@ -1,3 +1,7 @@
# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

from dataclasses import dataclass
from typing import Literal, Optional, Union

Expand Down
1 change: 1 addition & 0 deletions integrations/nvidia/tests/test_nim_backend.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import json
from unittest.mock import patch

Expand Down
6 changes: 0 additions & 6 deletions integrations/nvidia/tests/test_ranker.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,12 +46,6 @@ def test_init_pass_wo_api_key_w_api_url(self):
client = NvidiaRanker(api_url=url)
assert client.api_url == url

def test_warm_up_required(self):
client = NvidiaRanker()
with pytest.raises(RuntimeError) as e:
client.run("query", [Document(content="doc")])
assert "not been loaded" in str(e.value)

@pytest.mark.parametrize(
"truncate",
[
Expand Down
Loading