Commit 824785fb by 文靖昊

增加相关文档存储到excel的方法

parent caf42399
# -*- coding: utf-8 -*-
import sys
import openpyxl
import time
from langchain.chains import LLMChain
from langchain_core.prompts import PromptTemplate
......@@ -155,7 +156,48 @@ class QA:
self.cur_similarity = similarity.get_similarity_doc()
similarity_docs = similarity.get_similarity_docs()
rerank_docs = similarity.get_rerank_docs()
print(rerank_docs)
self.cur_question = self.prompt.format(history=history,context=self.cur_similarity, question=self.cur_oquestion)
if not _question:
return ""
self.cur_answer = self.llm.run(history=history,context=self.cur_similarity, question=self.cur_oquestion)
if self.contains_blocked_keywords(self.cur_answer):
self.cur_answer = SAFE_RESPONSE
self.update_history()
if with_similarity:
if self.rerank:
return self.cur_answer, rerank_docs
else:
return self.cur_answer, similarity_docs
return self.cur_answer
def chat_with_history_with_ext_with_save_excel(self, _question,ext,history,with_similarity=False):
self.cur_oquestion = _question
if self.contains_blocked_keywords(_question):
self.cur_answer = SAFE_RESPONSE
else:
# self.cur_similarity = self.get_similarity(_aquestion=self.cur_oquestion)
similarity1 = self.get_similarity_origin(_aquestion=self.cur_oquestion)
if self.rerank:
self.cur_similarity = similarity1.get_rerank(self.rerank_model)
else:
self.cur_similarity = similarity1.get_similarity_doc()
similarity_docs1 = similarity1.get_similarity_docs()
rerank_docs1 = similarity1.get_rerank_docs()
similarity = self.get_similarity_with_ext_origin(ext)
if self.rerank:
self.cur_similarity = similarity.get_rerank(self.rerank_model)
else:
self.cur_similarity = similarity.get_similarity_doc()
similarity_docs = similarity.get_similarity_docs()
rerank_docs = similarity.get_rerank_docs()
print("rerank_docs_with_ext================================================")
for i in range(len(rerank_docs)):
save_excel(_question,similarity_docs1[i].page_content, rerank_docs1[i].page_content,rerank_docs[i].page_content)
self.cur_question = self.prompt.format(history=history,context=self.cur_similarity, question=self.cur_oquestion)
if not _question:
return ""
......@@ -300,6 +342,28 @@ class QA:
n_info = info_llm.run(question=question)
self.crud.set_info(self.chat_id, n_info)
def save_excel(question,similar,rerank,rerank_ext):
file_full_path = r'D:\work\py\LAE\testdoc\test.xlsx'
# sheet名称
sheet_name = 'Sheet1'
# 获取指定的文件
wb = openpyxl.load_workbook(file_full_path)
# 获取指定的sheet
ws = wb[sheet_name]
# 获得最大行数
max_row_num = ws.max_row
# 获得最大列数
max_col_num = ws.max_column
ws.cell(row=max_row_num+1, column=1, value=question)
ws.cell(row=max_row_num+1, column=2, value=similar)
ws.cell(row=max_row_num+1, column=3, value=rerank)
ws.cell(row=max_row_num+1, column=4, value=rerank_ext)
# 保存文件
wb.save(file_full_path)
if __name__ == "__main__":
# 数据库连接
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment