Spaces:
Runtime error
Runtime error
File size: 2,762 Bytes
35b22df |
1 2 3 4 5 6 7 8 9 10 11 12 13 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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
"""Pinecone reader."""
from typing import Any, Dict, List, Optional
from gpt_index.readers.base import BaseReader
from gpt_index.readers.schema.base import Document
class PineconeReader(BaseReader):
"""Pinecone reader.
Args:
api_key (str): Pinecone API key.
environment (str): Pinecone environment.
"""
def __init__(self, api_key: str, environment: str):
"""Initialize with parameters."""
try:
import pinecone # noqa: F401
except ImportError:
raise ImportError(
"`pinecone` package not found, please run `pip install pinecone-client`"
)
self._api_key = api_key
self._environment = environment
pinecone.init(api_key=api_key, environment=environment)
def load_data(
self,
index_name: str,
id_to_text_map: Dict[str, str],
vector: Optional[List[float]],
top_k: int,
separate_documents: bool = True,
include_values: bool = True,
**query_kwargs: Any
) -> List[Document]:
"""Load data from Pinecone.
Args:
index_name (str): Name of the index.
id_to_text_map (Dict[str, str]): A map from ID's to text.
separate_documents (Optional[bool]): Whether to return separate
documents per retrieved entry. Defaults to True.
vector (List[float]): Query vector.
top_k (int): Number of results to return.
include_values (bool): Whether to include the embedding in the response.
Defaults to True.
**query_kwargs: Keyword arguments to pass to the query.
Arguments are the exact same as those found in
Pinecone's reference documentation for the
query method.
Returns:
List[Document]: A list of documents.
"""
import pinecone
index = pinecone.Index(index_name)
if "include_values" not in query_kwargs:
query_kwargs["include_values"] = True
response = index.query(top_k=top_k, vector=vector, **query_kwargs)
documents = []
for match in response.matches:
if match.id not in id_to_text_map:
raise ValueError("ID not found in id_to_text_map.")
text = id_to_text_map[match.id]
embedding = match.values
if len(embedding) == 0:
embedding = None
documents.append(Document(text=text, embedding=embedding))
if not separate_documents:
text_list = [doc.get_text() for doc in documents]
text = "\n\n".join(text_list)
documents = [Document(text=text)]
return documents
|