Spaces:
Sleeping
Sleeping
Create pinecone_utils.py
Browse files- pinecone_utils.py +34 -0
pinecone_utils.py
ADDED
|
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# pinecone_utils.py
|
| 2 |
+
|
| 3 |
+
import pinecone
|
| 4 |
+
from config import PINECONE_API_KEY, PINECONE_ENVIRONMENT, INDEX_NAME, CONTEXT_FIELDS
|
| 5 |
+
import torch
|
| 6 |
+
|
| 7 |
+
# Conectar a Pinecone
|
| 8 |
+
def connect_to_pinecone():
|
| 9 |
+
pinecone.init(api_key=PINECONE_API_KEY, environment=PINECONE_ENVIRONMENT)
|
| 10 |
+
index = pinecone.Index(INDEX_NAME)
|
| 11 |
+
return index
|
| 12 |
+
|
| 13 |
+
# Realizar búsqueda vectorial
|
| 14 |
+
def vector_search(query, embedding_model, index):
|
| 15 |
+
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
| 16 |
+
# Generar el embedding utilizando el modelo de embeddings
|
| 17 |
+
xq = embedding_model.encode(query, convert_to_tensor=True, device=device)
|
| 18 |
+
|
| 19 |
+
# Convertir el tensor a lista
|
| 20 |
+
xq = xq.cpu().tolist()
|
| 21 |
+
|
| 22 |
+
# Realizar búsqueda vectorial en el índice de Pinecone
|
| 23 |
+
res = index.query(vector=xq, top_k=3, include_metadata=True)
|
| 24 |
+
if res and res['matches']:
|
| 25 |
+
return [
|
| 26 |
+
{
|
| 27 |
+
'content': ' '.join(f"{k}: {v}" for k, v in match['metadata'].items() if k in CONTEXT_FIELDS and k != 'Tag'),
|
| 28 |
+
'metadata': match['metadata'],
|
| 29 |
+
'score': match.get('score', 0)
|
| 30 |
+
}
|
| 31 |
+
for match in res['matches']
|
| 32 |
+
if 'metadata' in match
|
| 33 |
+
]
|
| 34 |
+
return []
|