Spaces:
Sleeping
Sleeping
Update pinecone_utils.py
Browse files- pinecone_utils.py +25 -0
pinecone_utils.py
CHANGED
@@ -2,6 +2,8 @@
|
|
2 |
|
3 |
from pinecone import Pinecone, ServerlessSpec
|
4 |
from config import PINECONE_API_KEY, PINECONE_ENVIRONMENT, INDEX_NAME, CONTEXT_FIELDS
|
|
|
|
|
5 |
|
6 |
# Conectar a Pinecone
|
7 |
def connect_to_pinecone():
|
@@ -25,3 +27,26 @@ def connect_to_pinecone():
|
|
25 |
# Conectar al índice
|
26 |
index = pc.Index(INDEX_NAME)
|
27 |
return index
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
2 |
|
3 |
from pinecone import Pinecone, ServerlessSpec
|
4 |
from config import PINECONE_API_KEY, PINECONE_ENVIRONMENT, INDEX_NAME, CONTEXT_FIELDS
|
5 |
+
from sentence_transformers import SentenceTransformer
|
6 |
+
import torch
|
7 |
|
8 |
# Conectar a Pinecone
|
9 |
def connect_to_pinecone():
|
|
|
27 |
# Conectar al índice
|
28 |
index = pc.Index(INDEX_NAME)
|
29 |
return index
|
30 |
+
|
31 |
+
# Función para realizar la búsqueda vectorial
|
32 |
+
def vector_search(query, embedding_model, index):
|
33 |
+
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
34 |
+
# Generar el embedding utilizando el modelo de embeddings
|
35 |
+
xq = embedding_model.encode(query, convert_to_tensor=True, device=device)
|
36 |
+
|
37 |
+
# Convertir el tensor a lista
|
38 |
+
xq = xq.cpu().tolist()
|
39 |
+
|
40 |
+
# Realizar búsqueda vectorial en el índice de Pinecone
|
41 |
+
res = index.query(vector=xq, top_k=3, include_metadata=True)
|
42 |
+
if res and res['matches']:
|
43 |
+
return [
|
44 |
+
{
|
45 |
+
'content': ' '.join(f"{k}: {v}" for k, v in match['metadata'].items() if k in CONTEXT_FIELDS and k != 'Tag'),
|
46 |
+
'metadata': match['metadata'],
|
47 |
+
'score': match.get('score', 0)
|
48 |
+
}
|
49 |
+
for match in res['matches']
|
50 |
+
if 'metadata' in match
|
51 |
+
]
|
52 |
+
return []
|