blip / src /pinecone_store.py
amezi's picture
changing pinecone dims
2db5cdf
raw
history blame
1.04 kB
from pinecone import Pinecone, ServerlessSpec
import os
from dotenv import load_dotenv
load_dotenv()
class PineconeStore:
def __init__(self):
self.pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY"))
if 'blip-soccer-highlights-index' not in self.pc.list_indexes().names():
self.pc.create_index(
name='blip-soccer-highlights-index',
dimension=768, # This must match InternVL2_5-8B-MPO
metric='cosine',
spec=ServerlessSpec(cloud='aws', region=os.getenv("PINECONE_ENV", "us-east-1"))
)
self.index = self.pc.Index('blip-soccer-highlights-index')
def upsert(self, id, vector, metadata):
self.index.upsert([(id, vector.tolist(), metadata)])
def query(self, vector, filter_key, top_k):
results = self.index.query(vector.tolist(), top_k=top_k, include_metadata=True)
return [
match["metadata"] for match in results["matches"]
if filter_key in match["id"]
]