mirror of
https://github.com/hwchase17/langchain
synced 2024-11-04 06:00:26 +00:00
e6240fecab
- **Description:** Adds a RAG template that uses NVIDIA AI playground and embedding models, along with Milvus vector store - **Dependencies:** This template depends on the AI playground service in NVIDIA NGC. API keys with a significant trial compute are available (10k queries at the time of writing). This template also depends on the Milvus Vector store which is publicly available. Note: [A quick link to get a key](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/ai-foundation/models/codellama-13b/api) when you have an NGC account. Generate Key button at the top right of the code window. --------- Co-authored-by: Sagar B Manjunath <sbogadimanju@nvidia.com> Co-authored-by: Erick Friis <erick@langchain.dev>
40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
import getpass
|
|
import os
|
|
|
|
from langchain.document_loaders import PyPDFLoader
|
|
from langchain.text_splitter import CharacterTextSplitter
|
|
from langchain.vectorstores.milvus import Milvus
|
|
from langchain_nvidia_aiplay import NVIDIAEmbeddings
|
|
|
|
if os.environ.get("NVIDIA_API_KEY", "").startswith("nvapi-"):
|
|
print("Valid NVIDIA_API_KEY already in environment. Delete to reset")
|
|
else:
|
|
nvapi_key = getpass.getpass("NVAPI Key (starts with nvapi-): ")
|
|
assert nvapi_key.startswith("nvapi-"), f"{nvapi_key[:5]}... is not a valid key"
|
|
os.environ["NVIDIA_API_KEY"] = nvapi_key
|
|
|
|
# Note: if you change this, you should also change it in `nvidia_rag_canonical/chain.py`
|
|
EMBEDDING_MODEL = "nvolveqa_40k"
|
|
HOST = "127.0.0.1"
|
|
PORT = "19530"
|
|
COLLECTION_NAME = "test"
|
|
|
|
embeddings = NVIDIAEmbeddings(model=EMBEDDING_MODEL)
|
|
|
|
if __name__ == "__main__":
|
|
# Load docs
|
|
loader = PyPDFLoader("https://www.ssa.gov/news/press/factsheets/basicfact-alt.pdf")
|
|
data = loader.load()
|
|
|
|
# Split docs
|
|
text_splitter = CharacterTextSplitter(chunk_size=300, chunk_overlap=100)
|
|
docs = text_splitter.split_documents(data)
|
|
|
|
# Insert the documents in Milvus Vector Store
|
|
vector_db = Milvus.from_documents(
|
|
docs,
|
|
embeddings,
|
|
collection_name=COLLECTION_NAME,
|
|
connection_args={"host": HOST, "port": PORT},
|
|
)
|