76 lines
2.0 KiB
Python
76 lines
2.0 KiB
Python
import pickle
|
||
from graph import set_logger
|
||
from termcolor import colored
|
||
|
||
logger = set_logger(colored('BERT_VEC', 'yellow'))
|
||
bert_file_name = 'bert_data.pkl'
|
||
|
||
|
||
class BertData:
|
||
def __init__(self):
|
||
self.dic = {}
|
||
self._read_dic()
|
||
|
||
# 批量插入数据
|
||
def add_batch_data(self, keys, values):
|
||
for key, value in zip(keys, values):
|
||
self.dic[key] = value
|
||
|
||
# 插入单条数据
|
||
def add_data(self, key, value):
|
||
self.dic[key] = value
|
||
|
||
# 根据key删除数据
|
||
def delete_data(self, key):
|
||
if self.dic and self.dic.get(key, ''):
|
||
self.dic.pop(key)
|
||
|
||
# 根据key获取数据
|
||
def get_data(self, key):
|
||
return self.dic.get(key, '')
|
||
|
||
# 获取全部数据
|
||
def get_all_data(self):
|
||
return self.dic
|
||
|
||
# 提交
|
||
def commit(self):
|
||
self._save_dic()
|
||
|
||
def _save_dic(self):
|
||
try:
|
||
with open(bert_file_name, 'wb')as file:
|
||
pickle.dump(self.dic, file)
|
||
logger.info('bert data saved successfully')
|
||
except:
|
||
logger.info('save bert data failed')
|
||
|
||
def _read_dic(self):
|
||
try:
|
||
with open(bert_file_name, 'rb')as file:
|
||
self.dic = pickle.load(file)
|
||
except FileNotFoundError:
|
||
logger.info('local bert data is none')
|
||
|
||
|
||
if __name__ == '__main__':
|
||
bd = BertData()
|
||
data = []
|
||
vec = []
|
||
import numpy as np
|
||
|
||
for i in range(30000):
|
||
data.append('阿迪和考虑就鞍山市会计法哈三联空间和福利卡就很烦' + str(i))
|
||
vec.append(np.random.rand(768))
|
||
bd.add_batch_data(data, vec)
|
||
# 增删改需要调用commit方法才会修改本地缓存的内容，查询不需要调用该方法
|
||
bd.commit()
|
||
# bd.delete_data('上午好啊天气真的不错0')
|
||
# res = bd.get_data('上午好啊天气真的不错1')
|
||
# bd.add_data('上午好啊天气真的不错test', [1, 2, 3])
|
||
# res = bd.get_all_data()
|
||
# print(res.keys())
|
||
# for i in res.items():
|
||
# print(i[0], ':', i[1])
|
||
# print(res.values())
|