from intent_description_template import *
from FlagEmbedding import FlagAutoModel, FlagReranker
from intent_description_template import template
import torch
import os
import faiss
import numpy as np
import pandas as pd

num = 1
template_list = []
for k, v in template.items():
    for v2 in v:
        template_list.append((str(num), k, v2))
        num+=1
def print_gpu():
    if torch.cuda.is_available():
        print(f"allocated:{torch.cuda.memory_allocated()/1024**3:.2f}GB", end=' ')
        print(f"reserved: {torch.cuda.memory_reserved()/1024**3:.2f}GB")

# 加载模型
reranker = FlagReranker('./bge-reranker-v2-m3', use_fp16=True, local_files_only=True,devices=["cuda:0"])
# 模型预热
reranker.compute_score([("这是一段预热文字，首次推理通过预测保证后续推理的稳定性和性能。","这是一段预热文字，首次推理通过预测保证后续推理的稳定性和性能。")])
print_gpu()

# 开始计算矩阵
mat = np.zeros([len(template_list)+1, len(template_list)+1])
for i in range(len(template_list)+1):
    mat[i][0] = i
    mat[0][i] = i
for i in range(len(template_list)):
    query_template_pairs = []
    for j in range(len(template_list)):
        query_template_pairs.append((template_list[i][2], template_list[j][2]))
    relevant_score = reranker.compute_score(query_template_pairs, normalize=True)
    print(f"完成{i+1}行...")
    # 填充
    for j in range(len(template_list)):
        mat[i+1, j+1] = relevant_score[j]


# 将NumPy矩阵转换为DataFrame
df = pd.DataFrame(mat)
# 保存到Excel文件，不包含行索引和列标题
df.to_excel('template_similarity_matrix.xlsx', index=False, header=False, float_format='%.2f')

df = pd.DataFrame([("序号", "意图", "描述")] + template_list)
df.to_excel('template_description.xlsx', index=False, header=False)
pass