Spaces:
Runtime error
Runtime error
File size: 1,708 Bytes
e43be0f | 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 | from pinecone import Pinecone
from dotenv import load_dotenv
import os
load_dotenv()
class PineconeRetriever:
"""Class for initializing and querying Pinecone."""
def __init__(self, api_key, index_name, namespace):
self.api_key = api_key
self.index_name = index_name
self.namespace = namespace
self.pc = Pinecone(api_key=self.api_key)
self.index = self.pc.Index(self.index_name)
def retrieve_data(self, query, top_k=1):
"""Retrieve relevant data from Pinecone."""
# Convert the query into a numerical vector that Pinecone can search with
query_embedding = self.pc.inference.embed(
model="multilingual-e5-large",
inputs=[query],
parameters={
"input_type": "query",
"truncate": "END"
}
)
results = self.index.query(namespace=self.namespace,
vector=query_embedding[0].values,
top_k=top_k,
include_values=False,
include_metadata=True
)
return results
def retrieve_context(self, query, top_k=1):
results = self.retrieve_data(query, top_k)
context = ""
for match in results['matches']:
context += match['metadata']['source_text']
return context
if __name__ == "__main__":
pinecone_api=os.getenv("PINECONE_API_KEY"),
pinecone_index=os.getenv("PINECONE_INDEX"),
pinecone_namespace=os.getenv("PINECONE_NAMESPACE")
retriever = PineconeRetriever(API_KEY, INDEX_NAME, NAMESPACE)
query_result = retriever.retrieve_context("hi")
print(query_result)
|