import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

import numpy as np
import faiss
import json
from sentence_transformers import SentenceTransformer

d = 768  # dimension
zidonghua = np.load('zidonghua.npy')
model = SentenceTransformer('Dmeta-embedding-zh')

data = []
with open("data.json", encoding="utf-8") as f:
    for i in f.readlines():
        a = json.loads(i)
        data.append(a)

mubiaoliebie = "自动化技术"
data_prompt = []
for i in data:
    if str(i[1]) == "nan":
        continue

    leibie_list = i[1].split(";")
    for leibie in leibie_list:
        if leibie == mubiaoliebie:
            data_prompt.append("标题：“{}”，摘要：“{}”".format(i[0], i[2]))



# faiss.write_index(index, 'index.ivf')
index = faiss.read_index('zidonghua.ivf')

index.add(zidonghua)                    # add may be a bit slower as well
# D, I = index.search(xq, k)     # actual search
# print(I[-5:])                  # neighbors of the 5 last queries


print("=======================================")
index.nprobe = 2              # default nprobe is 1, try a few more
k = 4
biaoti = "工业机器人视觉导航系统的设计与实现"
zhaiyoa = "本研究致力于设计和实现工业机器人视觉导航系统，旨在提高工业生产中机器人的自主导航和定位能力。首先，通过综合考虑视觉传感器、定位算法和控制策略，设计了一种高效的机器人视觉导航系统框架。其次，利用深度学习技术对环境中的关键特征进行识别和定位，实现了机器人在复杂工作场景下的精确定位和路径规划。通过实验验证，本系统在提高机器人工作效率、减少人工干预以及降低操作误差等方面取得了显著的成果。因此，本研究为工业机器人在生产领域的应用提供了重要的技术支持，具有一定的实用和推广价值。"

prompt = "标题：“{}”，摘要：“{}”".format(biaoti, zhaiyoa)
embs = model.encode([prompt], normalize_embeddings=True)

D, I = index.search(embs, k)
print(I)

for i in I[0]:
    print(data_prompt[i])