Skip to content

nvidia_nim_rerank_driver

NvidiaNimRerankDriver

Bases: BaseRerankDriver

Nvidia Rerank Driver.

Source code in griptape/drivers/rerank/nvidia_nim_rerank_driver.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
@define(kw_only=True)
class NvidiaNimRerankDriver(BaseRerankDriver):
    """Nvidia Rerank Driver."""

    model: str = field()
    base_url: str = field()
    truncate: Literal["NONE", "END"] = field(default="NONE")
    headers: dict = field(factory=dict)

    def run(self, query: str, artifacts: list[TextArtifact]) -> list[TextArtifact]:
        if not artifacts:
            return []

        response = requests.post(
            url=f"{self.base_url.rstrip('/')}/v1/ranking",
            json=self._get_body(query, artifacts),
            headers=self.headers,
        )

        response.raise_for_status()

        ranked_artifacts = []
        for ranking in response.json()["rankings"]:
            artifact = artifacts[ranking["index"]]
            artifact.meta.update({"logit": ranking["logit"], "usage": ranking.get("usage")})
            ranked_artifacts.append(artifact)

        return ranked_artifacts

    def _get_body(self, query: str, artifacts: list[TextArtifact]) -> dict:
        return {
            "model": self.model,
            "query": {"text": query},
            "passages": [{"text": artifact.value} for artifact in artifacts],
            "truncate": self.truncate,
        }

base_url = field() class-attribute instance-attribute

headers = field(factory=dict) class-attribute instance-attribute

model = field() class-attribute instance-attribute

truncate = field(default='NONE') class-attribute instance-attribute

_get_body(query, artifacts)

Source code in griptape/drivers/rerank/nvidia_nim_rerank_driver.py 
43
44
45
46
47
48
49
def _get_body(self, query: str, artifacts: list[TextArtifact]) -> dict:
    return {
        "model": self.model,
        "query": {"text": query},
        "passages": [{"text": artifact.value} for artifact in artifacts],
        "truncate": self.truncate,
    }

run(query, artifacts)

Source code in griptape/drivers/rerank/nvidia_nim_rerank_driver.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def run(self, query: str, artifacts: list[TextArtifact]) -> list[TextArtifact]:
    if not artifacts:
        return []

    response = requests.post(
        url=f"{self.base_url.rstrip('/')}/v1/ranking",
        json=self._get_body(query, artifacts),
        headers=self.headers,
    )

    response.raise_for_status()

    ranked_artifacts = []
    for ranking in response.json()["rankings"]:
        artifact = artifacts[ranking["index"]]
        artifact.meta.update({"logit": ranking["logit"], "usage": ranking.get("usage")})
        ranked_artifacts.append(artifact)

    return ranked_artifacts