You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
128 lines
3.5 KiB
Python
128 lines
3.5 KiB
Python
|
|
import sys
|
|
|
|
from kb_config import logger
|
|
from sentence_transformers import CrossEncoder
|
|
from faiss_kb_service import FaissKBService, DocumentWithVectorStoreId
|
|
from langchain.docstore.document import Document
|
|
from base_kb import KnowledgeFile
|
|
|
|
class QAService():
|
|
def __init__(self, kb_name, device) -> None:
|
|
|
|
embed_model_path = 'bge-large-zh-v1.5'
|
|
fkbs = FaissKBService(kb_name, embed_model_path=embed_model_path, device=device)
|
|
fkbs.do_create_kb()
|
|
self.fkbs = fkbs
|
|
self.kb_name = kb_name
|
|
|
|
def delete_qa_file(self, qa_file_id):
|
|
kb_file = KnowledgeFile(qa_file_id, self.kb_name)
|
|
self.fkbs.do_delete_doc(kb_file, not_refresh_vs_cache=True)
|
|
|
|
def update_qa_doc(self, qa_file_id, doc_list, id_list):
|
|
self.delete_qa_file(qa_file_id)
|
|
|
|
doc_infos = self.fkbs.do_add_doc(doc_list, ids=id_list)
|
|
logger.info('fassi add docs: ' + str(len(doc_infos)))
|
|
|
|
self.fkbs.save_vector_store()
|
|
|
|
|
|
def search(self,
|
|
query,
|
|
top_k = 3,
|
|
score_threshold = 0.75,
|
|
reranked=False):
|
|
|
|
docs = self.fkbs.do_search(query, top_k, 1 - score_threshold)
|
|
|
|
return docs
|
|
|
|
|
|
import json
|
|
|
|
|
|
def create_question_id(intent_code, j, test_question):
|
|
return f"{intent_code}@{j}@{test_question}"
|
|
|
|
|
|
|
|
def load_testing_data(file_path):
|
|
|
|
test_data_list = []
|
|
question_list = []
|
|
id_list = []
|
|
|
|
with open(file_path, encoding='utf-8') as f:
|
|
data = json.load(f)
|
|
for i, item in enumerate(data):
|
|
test_question = item['testQuestion']
|
|
intent_code = item['expectIntentCode']
|
|
test_data_list.append((test_question, intent_code))
|
|
|
|
q_list = item['expectIntentQuestionExample']
|
|
for j, q in enumerate(q_list):
|
|
q_id = create_question_id(intent_code, j, test_question)
|
|
question_list.append(q)
|
|
id_list.append(q_id)
|
|
return test_data_list, question_list, id_list
|
|
|
|
|
|
def convert_to_doc_list(question_list, id_list, qa_file_id):
|
|
doc_list = []
|
|
for question, id in zip(question_list, id_list):
|
|
|
|
metadata = {
|
|
'source': qa_file_id,
|
|
'id': id
|
|
}
|
|
doc = Document(page_content=question, metadata=metadata)
|
|
doc_list.append(doc)
|
|
|
|
return doc_list
|
|
|
|
import time
|
|
def work():
|
|
start_time = time.time()
|
|
kb_name = 'my_kb_test'
|
|
device = None
|
|
qa_service = QAService(kb_name, device)
|
|
|
|
|
|
test_data_list, question_list, id_list = load_testing_data(r'test_data/testing_data.json')
|
|
print('Loaded data!')
|
|
|
|
qa_file_id = 'QA_TEST_2' # the source of the qa, using for data cleaning, make sure to be unique
|
|
|
|
|
|
doc_list = convert_to_doc_list(question_list, id_list, qa_file_id)
|
|
|
|
qa_service.update_qa_doc(qa_file_id, doc_list, id_list)
|
|
|
|
cnt = 0
|
|
for query, code in test_data_list:
|
|
rst = qa_service.search(query)
|
|
if do_test(query, code, rst):
|
|
cnt += 1
|
|
|
|
print(str(cnt) + '/' + str(len(test_data_list)))
|
|
|
|
end_time = time.time()
|
|
elapsed_time = end_time - start_time
|
|
print(f"总耗时: {elapsed_time} 秒")
|
|
|
|
def do_test(query, expected_intent_code, rst):
|
|
if rst is None or len(rst)==0:
|
|
print('Empty: ' + query)
|
|
return False
|
|
for rst_doc, similarity_score in rst:
|
|
page_content = rst_doc.page_content
|
|
intent_code = rst_doc.metadata['id'].split('@')[0]
|
|
print(
|
|
f"{query} vs {page_content} : {expected_intent_code} vs {intent_code} - Similarity Score: {1 - similarity_score}")
|
|
return True
|
|
|
|
work()
|
|
|