C2MV commited on
Commit
66bd678
1 Parent(s): 631287b

Update pinecone_utils.py

Browse files
Files changed (1) hide show
  1. 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 []