Spaces:
Runtime error
Runtime error
File size: 1,790 Bytes
2519bba |
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 |
import uuid
import pinecone
from pinecone import PineconeProtocolError
class Pinecone:
def __init__(self, api_key, environment, index_name):
self.api_key = api_key
self.environment = environment
self.index_name = index_name
self.index = None
def initialize_index(self):
pinecone.init(api_key=self.api_key, environment=self.environment)
self.index = pinecone.Index(self.index_name)
return self.index
def upsert_data(self, img_id, embeddings, path, img_url, page_url, face_coordinates, single_face):
vec_id = str(uuid.uuid4())
data = []
embedding_as_list = embeddings.tolist()
if face_coordinates is not None:
coordinates_1d = [str(coord) for coord in face_coordinates]
metadata = {'Image id': img_id, 'directory path': path, 'Image URL': img_url, 'Page URL': page_url, 'Face Coordinates': coordinates_1d, 'Single Face': single_face}
data.append((vec_id, embedding_as_list, metadata))
self.index.upsert(data)
def search_data(self, query_embedding):
try:
matches = self.index.query(
vector=query_embedding,
top_k=10,
include_values=True,
include_metadata = True
)
except PineconeProtocolError as e:
print(f"PineconeProtocolError occurred: {e}")
pinecone.deinit()
pinecone.init(api_key= self.api_key,environment=self.environment)
index = pinecone.Index(self.index_name)
matches = index.query(
vector=query_embedding,
top_k=10,
include_values=True,
include_metadata = True
)
return matches
|