llama-index postprocessor nvidia_rerank integration
Project description
LlamaIndex Postprocessor Integration: Nvidia_Rerank
The llama-index-postprocessor-nvidia-rerank
package contains LlamaIndex integrations for rerank model powered by the NVIDIA AI Foundation Model playground environment.
NVIDIA AI Foundation Endpoints give users easy access to hosted endpoints for generative AI models like Llama-2, SteerLM, Mistral, etc. Using the API, you can query live endpoints available on the NVIDIA GPU Cloud (NGC) to get quick results from a DGX-hosted cloud compute environment. All models are source-accessible and can be deployed on your own compute cluster. For more information about documentation, please visit NVIDIA NeMo Retriever Reranking.
Below is an example on how to use some common functionality surrounding text-generative and embedding models
Installation
pip install --upgrade llama-index llama-index-core llama-index-nvidia-rerank
Setup
To get started:
- Create a free account with the NVIDIA GPU Cloud service, which hosts AI solution catalogs, containers, models, etc.
- Navigate to
Catalog > AI Foundation Models > (Model with API endpoint)
. - Select the
API
option and clickGenerate Key
. - Save the generated key as
NVIDIA_API_KEY
. From there, you should have access to the endpoints.
This is how you set NVIDIA_API_KEY in environment variable export NVIDIA_API_KEY="Your_NVIDIA_API_KEY_obtained_from_above_setup"
import getpass
import os
if not os.environ.get("NVIDIA_API_KEY", "").startswith("nvapi-"):
nvidia_api_key = getpass.getpass("Enter your NVIDIA AIPLAY API key: ")
assert nvidia_api_key.startswith(
"nvapi-"
), f"{nvidia_api_key[:5]}... is not a valid key"
os.environ["NVIDIA_API_KEY"] = nvidia_api_key
from llama_index.postprocessor.nvidia_rerank import NVIDIARerank
# for local hosted nim exposed api end point
rerank = NVIDIARerank().mode(
mode="nim", base_url="http://<your_end_point>:1976/v1"
)
# for API Catalog reranker model
my_key = os.environ["NVIDIA_API_KEY"]
rerank = NVIDIARerank().mode(mode="nvidia", api_key=my_key)
Supported models
Querying get_available_models
will still give you all of the other models offered by your API credentials.
from llama_index.postprocessor.nvidia_rerank import NVIDIARerank
NVIDIARerank.get_available_models()
To find out more about a specific model, please navigate to the NVIDIA NIM section of ai.nvidia.com as linked here.
Reranking
Below is an example:
from llama_index.postprocessor.nvidia_rerank import NVIDIARerank
from llama_index.core import Document
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.node_parser import SentenceSplitter, SimpleFileNodeParser
import os
# load documents
documents = SimpleDirectoryReader("/path_to_your_data_folder").load_data()
# use API Catalog's reranker model
my_key = os.environ["NVIDIA_API_KEY"]
rerank = NVIDIARerank().mode(mode="nvidia", api_key=my_key)
# parse nodes
parser = SentenceSplitter(separator="\n", chunk_size=200, chunk_overlap=0)
nodes = parser.get_nodes_from_documents(documents)
# rerank
rerank.postprocess_nodes(nodes, query_str=query)
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for llama_index_postprocessor_nvidia_rerank-0.1.2.tar.gz
Algorithm | Hash digest | |
---|---|---|
SHA256 | 126ea3ce8ee647a65e35ef15c025cec7e075fd42fd601af73c90c36e9641e036 |
|
MD5 | 05b303b09bca300aedf9df996d7272df |
|
BLAKE2b-256 | 13495841417db5fb19b2e889dead58908304403fead3672d6b7a6b25237bbf28 |
Hashes for llama_index_postprocessor_nvidia_rerank-0.1.2-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | d67844923a80e27f2c4139c4f2147cc266a15d969e4a229e0999cb69ad9192e9 |
|
MD5 | bb243c512aa0c462a276dee72093b801 |
|
BLAKE2b-256 | fa99990d4f54636d344e8894b9fdf66caa39c8b7fca77a53505083a5bcddcf55 |