Dr.Yasuda_streamlit / api /pinecone_func.py
Blue2962
a
0207eb5
import os
from pinecone import Pinecone, ServerlessSpec
from openai import OpenAI
from dotenv import load_dotenv
load_dotenv()
pc = Pinecone(api_key=os.environ["PINECONE_API_KEY"])
if 'your-index' not in pc.list_indexes().names():
pc.create_index(
name='your-index',
dimension=1536,
metric='cosine',
spec=ServerlessSpec(
cloud='aws',
region='us-east-1'
)
)
client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
def upsert_texts(texts, ids):
embeddings = []
for text in texts:
response = client.embeddings.create(
model="text-embedding-ada-002",
input=text
)
embeddings.append(response.data[0].embedding)
pc = Pinecone(api_key=os.environ["PINECONE_API_KEY"])
index = pc.Index("your-index")
vectors = []
for id, emb, text in zip(ids, embeddings, texts):
vectors.append((id, emb, {"text": text}))
index.upsert(vectors)
def query_text(query, top_k=5):
pc = Pinecone(api_key=os.environ["PINECONE_API_KEY"])
index = pc.Index('your-index')
embedding = client.embeddings.create(
model="text-embedding-ada-002",
input=query
).data[0].embedding
results = index.query(vector=embedding, top_k=top_k, include_metadata=True)
return results