level_mode: 0: 原始节点 1: 聚合节点 2: 所有节点
(working_dir,query,topk=10,level_mode=2)
| 66 | # ) |
| 67 | |
| 68 | def search_vector_search(working_dir,query,topk=10,level_mode=2): |
| 69 | ''' |
| 70 | level_mode: 0: 原始节点 |
| 71 | 1: 聚合节点 |
| 72 | 2: 所有节点 |
| 73 | ''' |
| 74 | if level_mode==0: |
| 75 | filter_filed=" level == 0 " |
| 76 | elif level_mode==1: |
| 77 | filter_filed=" level > 0 " |
| 78 | # elif level_mode==2: |
| 79 | # filter_filed=" level < 58736" |
| 80 | else: |
| 81 | filter_filed="" |
| 82 | dataset=os.path.basename(working_dir) |
| 83 | if os.path.exists(f"{working_dir}/milvus_demo.db"): |
| 84 | print(f"{working_dir}milvus_demo.db already exists, using it") |
| 85 | milvus_client = MilvusClient(uri=f"{working_dir}/milvus_demo.db") |
| 86 | else: |
| 87 | print("milvus_demo.db not found, using default") |
| 88 | milvus_client = MilvusClient(uri=f"/data/zyz/trag_ds/exp/ds_hire_cs20_top20_chunk5/{dataset}/milvus_demo.db") |
| 89 | collection_name = "entity_collection" |
| 90 | # query_embedding = emb_text(query) |
| 91 | search_results = milvus_client.search( |
| 92 | collection_name=collection_name, |
| 93 | data=query, |
| 94 | limit=topk, |
| 95 | params={"metric_type": "IP", "params": {}}, |
| 96 | filter=filter_filed, |
| 97 | output_fields=["entity_name", "description","parent","level","source_id"], |
| 98 | ) |
| 99 | # print(search_results) |
| 100 | extract_results=[(i['entity']['entity_name'],i["entity"]["parent"],i["entity"]["description"],i["entity"]["source_id"])for i in search_results[0]] |
| 101 | # print(extract_results) |
| 102 | return extract_results |
| 103 | def create_db_table_mysql(working_dir): |
| 104 | con = pymysql.connect(host='localhost',port=4321, user='root', |
| 105 | passwd='123', charset='utf8mb4') |