Skip to content

Rerank

Overview

Rerank Drivers can be used to rerank search results for a particular query. Every Rerank Driver implements the following methods that can be used directly:

  • run(query: str, artifacts: list[TextArtifact]) reranks a list of TextArtifact based on the original query.

Rerank Drivers can also be used with a RagEngine's Rerank Module.

Rerank Drivers

Local

The LocalRerankDriver uses a simple relatedness calculation.

from griptape.artifacts import TextArtifact
from griptape.drivers.rerank.local import LocalRerankDriver

rerank_driver = LocalRerankDriver()

artifacts = rerank_driver.run(
    "What is the capital of France?",
    [
        TextArtifact("Hotdog"),
        TextArtifact("San Francisco"),
        TextArtifact("Paris"),
        TextArtifact("Baguette"),
        TextArtifact("French Hotdog"),
    ],
)
for artifact in artifacts:
    print(artifact.value)

Paris
San Francisco
Baguette
French Hotdog
Hotdog

Cohere

The CohereRerankDriver uses Cohere's Rerank model.

import os

from griptape.artifacts import TextArtifact
from griptape.drivers.rerank.cohere import CohereRerankDriver

rerank_driver = CohereRerankDriver(
    api_key=os.environ["COHERE_API_KEY"],
)

artifacts = rerank_driver.run(
    "Where is NYC located?",
    [
        TextArtifact("NYC Media"),
        TextArtifact("New York City Police Department"),
        TextArtifact("New York City"),
        TextArtifact("New York City Subway"),
    ],
)
for artifact in artifacts:
    print(artifact.value)

New York City
NYC Media
New York City Police Department
New York City Subway