# Install these packages:
# !pip install pinecone
# !pip install pinecone_datasets

from pinecone import Pinecone, ServerlessSpec
from pinecone_datasets import load_dataset

# Initialize Pinecone
pc = Pinecone(api_key="YOUR_API_KEY")

# Create the index if it doesn't already exist
pc.create_index(name="hello-pinecone", dimension=1536, spec=ServerlessSpec(cloud='aws', region='us-east-1'))

# Connect to the index
index = pc.Index(name='hello-pinecone')

# Load the dataset
dataset = load_dataset("langchain-python-docs-text-embedding-ada-002")

# Define a function to upsert batches
def upsert_batch(batch):
    index.upsert(vectors=batch)

print("Your index has been created! Uploading a few items from the dataset...")

# Upsert the first 1000 entries from the dataset
for i, batch in enumerate(dataset.iter_documents(batch_size=50)):
    index.upsert(vectors=batch)
    if (i + 1) * 50 >= 1000:
        break

# Function to get the first entry
def get_first_entry(dataset):
    for batch in dataset.iter_documents(batch_size=1):
        return batch[0]['values'].tolist()

# Query the index
print(index.query(
    vector=get_first_entry(dataset),
    top_k=1,
    include_values=True,
    include_metadata=True
))