-
Notifications
You must be signed in to change notification settings - Fork 0
/
queryExample.py
33 lines (24 loc) · 871 Bytes
/
queryExample.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
from sentence_transformers import SentenceTransformer
import torch
import pinecone
import settings
import utils
PINECONE_API_KEY = settings.PINECONE_API_KEY
PINECONE_ENV = settings.PINECONE_ENV
pinecone.init(
api_key=PINECONE_API_KEY,
environment=PINECONE_ENV
)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = SentenceTransformer('all-MiniLM-L6-v2', device=device)
model
# Enter desired query here
query = "Foreign trade sanctions"
# create the query vector
xq = model.encode(query).tolist()
index_name = "semantic-pdf"
index = pinecone.GRPCIndex(index_name)
xc = index.query(xq, top_k=5, include_metadata=True)
for result in xc['matches']:
summary = utils.get_summary(result['metadata']['text'])
print(f"{round(result['score'], 2)} Match - id: {result['id']}, pdf_file: {result['metadata']['pdf_file']}, text summary: {summary}")