|
2
|
1 import json
|
|
|
2 from pathlib import Path
|
|
|
3
|
|
11
|
4 # from pymilvus import MilvusClient
|
|
|
5 # from milvus_model.dense.onnx import OnnxEmbeddingFunction
|
|
|
6 from fastapi import FastAPI
|
|
2
|
7 from tqdm import tqdm
|
|
|
8
|
|
|
9
|
|
|
10 def rebuild(client, embedding_fn, dim):
|
|
|
11 client.drop_collection(collection_name="demo_collection")
|
|
|
12 if not client.has_collection(collection_name="demo_collection"):
|
|
|
13 client.create_collection(
|
|
|
14 collection_name="demo_collection",
|
|
|
15 dimension=dim,
|
|
|
16 )
|
|
|
17
|
|
|
18 docs = []
|
|
|
19 for i, (bbox, phrase) in tqdm(enumerate(
|
|
|
20 phrasesFromFile(
|
|
|
21 Path("data") /
|
|
4
|
22 "Meetings2226Minutes_20240702182359526 (1).pdf")),
|
|
8
|
23 desc="rebuilding",
|
|
|
24 unit=' phrase'):
|
|
2
|
25 [vector] = embedding_fn.encode_documents([phrase])
|
|
|
26 doc = {
|
|
|
27 "id": i,
|
|
|
28 "vector": vector,
|
|
|
29 "text": phrase,
|
|
|
30 "bbox": json.dumps(bbox),
|
|
|
31 }
|
|
|
32 docs.append(doc)
|
|
|
33 res = client.insert(collection_name="demo_collection", data=docs)
|
|
4
|
34 print('insert:', res['insert_count'])
|
|
2
|
35
|
|
|
36
|
|
8
|
37 def xxsearch(q, embedding_fn, client):
|
|
2
|
38 query_vectors = embedding_fn.encode_queries([q])
|
|
|
39
|
|
|
40 [query_result] = client.search(
|
|
|
41 collection_name="demo_collection",
|
|
|
42 data=query_vectors,
|
|
|
43 limit=5,
|
|
|
44 output_fields=["text"],
|
|
|
45 )
|
|
|
46 query_result.sort(key=lambda x: x["distance"], reverse=True)
|
|
|
47
|
|
|
48 for row in query_result:
|
|
|
49 print(f'{row["distance"]:.6f} {row["entity"]["text"]}')
|
|
|
50
|
|
0
|
51
|
|
4
|
52 # q, = sys.argv[1:]
|
|
|
53
|
|
|
54 # https://huggingface.co/models?pipeline_tag=feature-extraction&library=onnx&sort=trending
|
|
|
55 # embedding_fn = OnnxEmbeddingFunction(model_name="jinaai/jina-embeddings-v2-base-en")
|
|
|
56 # client = MilvusClient("milvus_demo.db")
|
|
|
57 # rebuild(client, embedding_fn, dim=embedding_fn.dim)
|
|
|
58 # search(q, embedding_fn, client)
|
|
0
|
59
|
|
4
|
60 app = FastAPI()
|
|
|
61
|
|
11
|
62 #search = Search()
|
|
8
|
63
|
|
4
|
64
|
|
|
65 @app.get("/sco/query")
|
|
8
|
66 def read_query1(q: str):
|
|
|
67 results = []
|
|
|
68 results = search.search(q)
|
|
|
69
|
|
|
70 return {"results": results}
|
