import numpy as np
from sentence_transformers import SentenceTransformer


encoder = SentenceTransformer("../model/")

tags = []
f = open('tags.txt', 'r')
for line in f.readlines():
    tags.append(line.strip())
f.close()

tags_embed = encoder.encode(tags)
tags_dis = [np.sqrt(np.dot(_, _.T)) for _ in tags_embed]

print(tags_embed.shape, tags_dis.shape)

with open('./tags_embed.npy', 'wb') as f:
    np.save(f, tags_embed)

with open('./tags_dis.npy', 'wb') as f:
    np.save(f, tags_dis)

# f = open('gpttag.txt', 'r')
# data = eval(f.readline())
# f.close()
#
# out = ""
# f = open('tags.txt', 'w')
# for tmp in data:
#     if tmp[1] > 2:
#         out += tmp[0] + '\n'
#
# f.write(out)
# f.close()