|
|
|
|
|
# coding=utf-8
|
|
|
|
|
|
from 数据处理 import embed_data,dBgx
|
|
|
|
|
|
import random as rand
|
|
|
|
|
|
from langchain import PromptTemplate, LLMChain
|
|
|
|
|
|
from 模型 import CustomLLM
|
|
|
|
|
|
from args import args
|
|
|
|
|
|
import torch
|
|
|
|
|
|
from transformers import AutoTokenizer, AutoModel,AutoConfig
|
|
|
|
|
|
from langchain.text_splitter import CharacterTextSplitter
|
|
|
|
|
|
import re
|
|
|
|
|
|
from typing import List
|
|
|
|
|
|
from langchain.document_loaders.unstructured import UnstructuredFileLoader
|
|
|
|
|
|
from paddleocr import PaddleOCR
|
|
|
|
|
|
import os
|
|
|
|
|
|
import fitz
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
zhishiku_init1='''欢迎来到基于知识库的菜单推荐，模型的回答效果除了与模型能力有关外，还需注意创建知识库时根据数据选择合适的切分长度和重叠长度，喜好提交时尽量详细数据库里存在的内容'''
|
|
|
|
|
|
caidan_init0="""欢迎来到基于菜单的菜单推荐，询问前请先上传菜单"""
|
|
|
|
|
|
damodel_init0="""欢迎来到大模型随机推荐，输入你的喜好，模型将根据你的喜好随机推荐"""
|
|
|
|
|
|
zhinengshibie = '''欢迎使用智能识别，使用前请先去配置文件填写相关的api。
|
|
|
|
|
|
api申请方法见：https://blog.csdn.net/2303_79001442/article/details/132093208
|
|
|
|
|
|
别问为什么不用本地模型，问就是显存不足。显存足够的可以考虑更换自己的模型。'''
|
|
|
|
|
|
def get_answer(query,chatbot):
|
|
|
|
|
|
|
|
|
|
|
|
llm = CustomLLM()
|
|
|
|
|
|
db = embed_data(args.embeddings_model_name,
|
|
|
|
|
|
args.original_data_path,
|
|
|
|
|
|
args.preprocessed_data_path,
|
|
|
|
|
|
)
|
|
|
|
|
|
retriever = db.as_retriever(search_kwargs={"k": args.topk_1})
|
|
|
|
|
|
docs = retriever.get_relevant_documents(query)
|
|
|
|
|
|
prompt_template = """根据已知内容推荐几盘菜肴。已知内容: {adjective}"""
|
|
|
|
|
|
|
|
|
|
|
|
prompt = PromptTemplate(template=prompt_template, input_variables=["adjective"])
|
|
|
|
|
|
|
|
|
|
|
|
llm_chain = LLMChain(prompt=prompt, llm=llm)
|
|
|
|
|
|
a = llm_chain.run(adjective=docs)
|
|
|
|
|
|
|
|
|
|
|
|
b=[query,a]
|
|
|
|
|
|
chatbot.append(b)
|
|
|
|
|
|
|
|
|
|
|
|
return chatbot, ''
|
|
|
|
|
|
from langchain.vectorstores import FAISS
|
|
|
|
|
|
from langchain.embeddings import HuggingFaceEmbeddings
|
|
|
|
|
|
def get_answer_caidan(query, chatbot):
|
|
|
|
|
|
embeddings = HuggingFaceEmbeddings(model_name=args.embeddings_model_name)
|
|
|
|
|
|
llm = CustomLLM()
|
|
|
|
|
|
caidan = load_caidan(args.path_caidan,sentence_size=args.chunk_size) # 加载菜单
|
|
|
|
|
|
db = FAISS.from_documents(caidan, embeddings)
|
|
|
|
|
|
retriever = db.as_retriever(search_kwargs={"k": args.topk_1})
|
|
|
|
|
|
docs = retriever.get_relevant_documents(query)
|
|
|
|
|
|
|
|
|
|
|
|
prompt = """基于我提供的菜单信息和饮食喜好信息。简洁的回答我的问题。
|
|
|
|
|
|
菜单信息：{adjective}
|
|
|
|
|
|
饮食喜好信息：{xihao}
|
|
|
|
|
|
问题:{question}"""
|
|
|
|
|
|
question = '''请根据我的喜好从菜单中选出我可能喜欢的菜肴，答案中不允许出现菜单中没有的信息。'''
|
|
|
|
|
|
|
|
|
|
|
|
prompt = PromptTemplate(template=prompt, input_variables=["adjective", "xihao", "question"])
|
|
|
|
|
|
llm_chain = LLMChain(prompt=prompt, llm=llm)
|
|
|
|
|
|
a = llm_chain.run(adjective=docs, xihao=query, question=question)
|
|
|
|
|
|
b = [query, a]
|
|
|
|
|
|
chatbot.append(b)
|
|
|
|
|
|
|
|
|
|
|
|
return chatbot, ''
|
|
|
|
|
|
def get_answer_model(query, chatbot):
|
|
|
|
|
|
|
|
|
|
|
|
llm = CustomLLM()
|
|
|
|
|
|
prompt = """你现在是一位优秀的厨师，你要根据我的饮食喜好推荐几种美食。我的饮食喜好如下：{xihao}。回答只能展示美食的名字不需要过多的解释。"""
|
|
|
|
|
|
prompt = PromptTemplate(template=prompt, input_variables=["xihao"])
|
|
|
|
|
|
llm_chain = LLMChain(prompt=prompt, llm=llm)
|
|
|
|
|
|
a = llm_chain.run(xihao=query)
|
|
|
|
|
|
b = [query, a]
|
|
|
|
|
|
chatbot.append(b)
|
|
|
|
|
|
|
|
|
|
|
|
return chatbot, ''
|
|
|
|
|
|
def get_answer_guihua(query, chatbot,mode2):
|
|
|
|
|
|
llm = CustomLLM()
|
|
|
|
|
|
if mode2 == '启用':
|
|
|
|
|
|
db = embed_data(args.embeddings_model_name,
|
|
|
|
|
|
args.original_data_path_g,
|
|
|
|
|
|
args.preprocessed_data_path_g,
|
|
|
|
|
|
)
|
|
|
|
|
|
retriever = db.as_retriever(search_kwargs={"k": args.topk_1})
|
|
|
|
|
|
docs = retriever.get_relevant_documents(query)
|
|
|
|
|
|
prompt_template = """请参考以下的素材信息和已有食材为我推荐几种菜肴，并指出缺少的食材和具体的操作过程。
|
|
|
|
|
|
素材信息：{adjective}，
|
|
|
|
|
|
已有食材：{shicai}"""
|
|
|
|
|
|
|
|
|
|
|
|
prompt = PromptTemplate(template=prompt_template, input_variables=["adjective",'shicai'])
|
|
|
|
|
|
|
|
|
|
|
|
llm_chain = LLMChain(prompt=prompt, llm=llm)
|
|
|
|
|
|
a = llm_chain.run(adjective=docs,shicai=query)
|
|
|
|
|
|
|
|
|
|
|
|
b = [query, a]
|
|
|
|
|
|
chatbot.append(b)
|
|
|
|
|
|
|
|
|
|
|
|
elif mode2 == '不启用':
|
|
|
|
|
|
prompt = """你是一位优秀的厨师，我将提供一些食材，你要根据食材制作几种不同的美味菜肴。提供的食材如下：{食材}。"""
|
|
|
|
|
|
prompt = PromptTemplate(template=prompt, input_variables=["食材"])
|
|
|
|
|
|
llm_chain = LLMChain(prompt=prompt, llm=llm)
|
|
|
|
|
|
a = llm_chain.run(食材=query)
|
|
|
|
|
|
b = [query, a]
|
|
|
|
|
|
chatbot.append(b)
|
|
|
|
|
|
|
|
|
|
|
|
return chatbot, ''
|
|
|
|
|
|
def get_answer_guihua2(query, chatbot,mode2):
|
|
|
|
|
|
llm = CustomLLM()
|
|
|
|
|
|
if mode2 == '启用':
|
|
|
|
|
|
db = embed_data(args.embeddings_model_name,
|
|
|
|
|
|
args.original_data_path_g,
|
|
|
|
|
|
args.preprocessed_data_path_g)
|
|
|
|
|
|
retriever = db.as_retriever(search_kwargs={"k": args.topk_1})
|
|
|
|
|
|
docs = retriever.get_relevant_documents(query)
|
|
|
|
|
|
prompt_template = """请参考以下的素材信息回答{shicai}制作过程。
|
|
|
|
|
|
素材信息：{adjective}"""
|
|
|
|
|
|
|
|
|
|
|
|
prompt = PromptTemplate(template=prompt_template, input_variables=["adjective",'shicai'])
|
|
|
|
|
|
|
|
|
|
|
|
llm_chain = LLMChain(prompt=prompt, llm=llm)
|
|
|
|
|
|
a = llm_chain.run(adjective=docs,shicai=query)
|
|
|
|
|
|
|
|
|
|
|
|
b = [query, a]
|
|
|
|
|
|
chatbot.append(b)
|
|
|
|
|
|
|
|
|
|
|
|
elif mode2 == '不启用':
|
|
|
|
|
|
prompt = """你是一位优秀的厨师，我想制作一份{shicai}，请问具体应该怎么做。"""
|
|
|
|
|
|
prompt = PromptTemplate(template=prompt, input_variables=["shicai"])
|
|
|
|
|
|
llm_chain = LLMChain(prompt=prompt, llm=llm)
|
|
|
|
|
|
a = llm_chain.run(shicai=query)
|
|
|
|
|
|
b = [query, a]
|
|
|
|
|
|
chatbot.append(b)
|
|
|
|
|
|
|
|
|
|
|
|
return chatbot, ''
|
|
|
|
|
|
|
|
|
|
|
|
def ziyouduihau(query, chatbot):
|
|
|
|
|
|
payload = {
|
|
|
|
|
|
"temperature": args.temperature,
|
|
|
|
|
|
"top_k": args.topk,
|
|
|
|
|
|
}
|
|
|
|
|
|
model_config = AutoConfig.from_pretrained(args.pretrained_model_name, trust_remote_code=True)
|
|
|
|
|
|
model_config.update(payload)
|
|
|
|
|
|
|
|
|
|
|
|
model = AutoModel.from_pretrained(args.pretrained_model_name, config=model_config,trust_remote_code=True).half().to(args.device)
|
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(args.pretrained_model_name, trust_remote_code=True)
|
|
|
|
|
|
model = model.eval()
|
|
|
|
|
|
prompt = query
|
|
|
|
|
|
response, history = model.chat(tokenizer, prompt, history=[])
|
|
|
|
|
|
b = [query, response]
|
|
|
|
|
|
chatbot.append(b)
|
|
|
|
|
|
|
|
|
|
|
|
return chatbot, ''
|
|
|
|
|
|
def change_data_path(data_path):
|
|
|
|
|
|
args.original_data_path = data_path
|
|
|
|
|
|
def change_data_path_g(data_path):
|
|
|
|
|
|
args.original_data_path = data_path
|
|
|
|
|
|
def change_P_data_path(data_path):
|
|
|
|
|
|
args.preprocessed_data_path = data_path
|
|
|
|
|
|
def change_P_data_path_g(data_path):
|
|
|
|
|
|
args.preprocessed_data_path_g = data_path
|
|
|
|
|
|
def change_emb_name(emb_name):
|
|
|
|
|
|
args.embeddings_model_name = emb_name
|
|
|
|
|
|
def change_search_top_k(search_top_k):
|
|
|
|
|
|
if search_top_k != '':
|
|
|
|
|
|
args.topk_1 = int(search_top_k)
|
|
|
|
|
|
def change_model_top_k(model_top_k):
|
|
|
|
|
|
if model_top_k != '':
|
|
|
|
|
|
args.topk = int(model_top_k)
|
|
|
|
|
|
def change_model_temperature(model_temperature):
|
|
|
|
|
|
args.temperature = float(model_temperature)
|
|
|
|
|
|
def change_model_name(model_name):
|
|
|
|
|
|
args.pretrained_model_name = model_name
|
|
|
|
|
|
def db_update_click1(chatbot):
|
|
|
|
|
|
|
|
|
|
|
|
embeddings_model_name =args.embeddings_model_name
|
|
|
|
|
|
|
|
|
|
|
|
original_data_path = args.original_data_path
|
|
|
|
|
|
preprocessed_data_path = args.preprocessed_data_path
|
|
|
|
|
|
|
|
|
|
|
|
dBgx(original_data_path,preprocessed_data_path,embeddings_model_name)
|
|
|
|
|
|
|
|
|
|
|
|
dbok = [None, '知识库已更新']
|
|
|
|
|
|
chatbot.append(dbok)
|
|
|
|
|
|
return chatbot
|
|
|
|
|
|
def db_update_click(chatbot):
|
|
|
|
|
|
|
|
|
|
|
|
embeddings_model_name = args.embeddings_model_name
|
|
|
|
|
|
|
|
|
|
|
|
original_data_path = args.original_data_path_g
|
|
|
|
|
|
preprocessed_data_path = args.preprocessed_data_path_g
|
|
|
|
|
|
|
|
|
|
|
|
dBgx(original_data_path,preprocessed_data_path,embeddings_model_name)
|
|
|
|
|
|
dbok = [None, '知识库已更新']
|
|
|
|
|
|
chatbot.append(dbok)
|
|
|
|
|
|
return chatbot
|
|
|
|
|
|
def change_chunk_size(chunk_size):
|
|
|
|
|
|
args.chunk_size = chunk_size
|
|
|
|
|
|
def change_chunk_overlap(chunk_overlap):
|
|
|
|
|
|
args.chunk_overlap = chunk_overlap
|
|
|
|
|
|
def clean_memory(chatbot):
|
|
|
|
|
|
if torch.cuda.is_available():
|
|
|
|
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
torch.cuda.ipc_collect()
|
|
|
|
|
|
|
|
|
|
|
|
elif torch.backends.mps.is_available():
|
|
|
|
|
|
try:
|
|
|
|
|
|
from torch.mps import empty_cache
|
|
|
|
|
|
empty_cache()
|
|
|
|
|
|
except Exception as e:
|
|
|
|
|
|
print(e)
|
|
|
|
|
|
print("如果您使用的是 macOS 建议将 pytorch 版本升级至 2.0.0 或更高版本，以支持及时清理 torch 产生的内存占用。")
|
|
|
|
|
|
clmok = [None, '显存已清理']
|
|
|
|
|
|
chatbot.append(clmok)
|
|
|
|
|
|
|
|
|
|
|
|
return chatbot
|
|
|
|
|
|
def change_caidan_path(caidan):
|
|
|
|
|
|
args.path_caidan = caidan
|
|
|
|
|
|
def change_path_zidingyi(suijidz):
|
|
|
|
|
|
args.path_zidingyi = suijidz
|
|
|
|
|
|
def change_suiji_num(suiji_num):
|
|
|
|
|
|
args.num_select = suiji_num
|
|
|
|
|
|
def change_shitu_top_k(shitu_top_k):
|
|
|
|
|
|
args.top_num = shitu_top_k
|
|
|
|
|
|
def change_shitu_filter_threshold(shitu_filter_threshold):
|
|
|
|
|
|
args.filter_threshold = shitu_filter_threshold
|
|
|
|
|
|
def change_shitu_baike_num(shitu_baike_num):
|
|
|
|
|
|
args.baike_num = shitu_baike_num
|
|
|
|
|
|
def Random_selection1(chatbot):
|
|
|
|
|
|
content = []
|
|
|
|
|
|
with open(args.path_zidingyi, 'r',encoding='utf-8') as file:
|
|
|
|
|
|
for line in file:
|
|
|
|
|
|
content.append(line.strip())
|
|
|
|
|
|
data_len = len(content)
|
|
|
|
|
|
num_select = args.num_select
|
|
|
|
|
|
random_rows = rand.sample(range(data_len), num_select)
|
|
|
|
|
|
elements = [content[i] for i in random_rows]
|
|
|
|
|
|
shuchu='随机内容为:\n' + '\n '.join(elements)
|
|
|
|
|
|
b = [None, shuchu]
|
|
|
|
|
|
chatbot.append(b)
|
|
|
|
|
|
return chatbot
|
|
|
|
|
|
def Random_selection2(chatbot,content_list):
|
|
|
|
|
|
|
|
|
|
|
|
content_list = content_list.split("\n")
|
|
|
|
|
|
data_len = len(content_list)
|
|
|
|
|
|
num_select = args.num_select
|
|
|
|
|
|
random_rows = rand.sample(range(data_len), num_select)
|
|
|
|
|
|
elements = [content_list[i] for i in random_rows]
|
|
|
|
|
|
shuchu='随机内容为:\n' + '\n '.join(elements)
|
|
|
|
|
|
b = [None, shuchu]
|
|
|
|
|
|
chatbot.append(b)
|
|
|
|
|
|
|
|
|
|
|
|
return chatbot
|
|
|
|
|
|
class UnstructuredPaddleImageLoader(UnstructuredFileLoader):
|
|
|
|
|
|
"""Loader that uses unstructured to load image files, such as PNGs and JPGs."""
|
|
|
|
|
|
|
|
|
|
|
|
def _get_elements(self) -> List:
|
|
|
|
|
|
def image_ocr_txt(filepath, dir_path="tmp_files"):
|
|
|
|
|
|
full_dir_path = os.path.join(os.path.dirname(filepath), dir_path)
|
|
|
|
|
|
if not os.path.exists(full_dir_path):
|
|
|
|
|
|
os.makedirs(full_dir_path)
|
|
|
|
|
|
filename = os.path.split(filepath)[-1]
|
|
|
|
|
|
ocr = PaddleOCR(use_angle_cls=True, lang="ch", use_gpu=False, show_log=False)
|
|
|
|
|
|
result = ocr.ocr(img=filepath)
|
|
|
|
|
|
|
|
|
|
|
|
ocr_result = [i[1][0] for line in result for i in line]
|
|
|
|
|
|
txt_file_path = os.path.join(full_dir_path, "%s.txt" % (filename))
|
|
|
|
|
|
with open(txt_file_path, 'w', encoding='utf-8') as fout:
|
|
|
|
|
|
fout.write("\n".join(ocr_result))
|
|
|
|
|
|
return txt_file_path
|
|
|
|
|
|
|
|
|
|
|
|
txt_file_path = image_ocr_txt(self.file_path)
|
|
|
|
|
|
from unstructured.partition.text import partition_text
|
|
|
|
|
|
return partition_text(filename=txt_file_path, **self.unstructured_kwargs)
|
|
|
|
|
|
|
|
|
|
|
|
class ChineseTextSplitter(CharacterTextSplitter):
|
|
|
|
|
|
def __init__(self, pdf: bool = False, sentence_size: int = 100, **kwargs):
|
|
|
|
|
|
super().__init__(**kwargs)
|
|
|
|
|
|
self.pdf = pdf
|
|
|
|
|
|
self.sentence_size = sentence_size
|
|
|
|
|
|
|
|
|
|
|
|
def split_text1(self, text: str) -> List[str]:
|
|
|
|
|
|
if self.pdf:
|
|
|
|
|
|
text = re.sub(r"\n{3,}", "\n", text)
|
|
|
|
|
|
text = re.sub('\s', ' ', text)
|
|
|
|
|
|
text = text.replace("\n\n", "")
|
|
|
|
|
|
sent_sep_pattern = re.compile('([﹒﹔﹖﹗．。！？]["’”」』]{0,2}|(?=["‘“「『]{1,2}|$))') # del ：；
|
|
|
|
|
|
sent_list = []
|
|
|
|
|
|
for ele in sent_sep_pattern.split(text):
|
|
|
|
|
|
if sent_sep_pattern.match(ele) and sent_list:
|
|
|
|
|
|
sent_list[-1] += ele
|
|
|
|
|
|
elif ele:
|
|
|
|
|
|
sent_list.append(ele)
|
|
|
|
|
|
return sent_list
|
|
|
|
|
|
|
|
|
|
|
|
def split_text(self, text: str) -> List[str]: ##此处需要进一步优化逻辑
|
|
|
|
|
|
if self.pdf:
|
|
|
|
|
|
text = re.sub(r"\n{3,}", r"\n", text)
|
|
|
|
|
|
text = re.sub('\s', " ", text)
|
|
|
|
|
|
text = re.sub("\n\n", "", text)
|
|
|
|
|
|
|
|
|
|
|
|
text = re.sub(r'([;；.!?。！？\?])([^”’])', r"\1\n\2", text) # 单字符断句符
|
|
|
|
|
|
text = re.sub(r'(\.{6})([^"’”」』])', r"\1\n\2", text) # 英文省略号
|
|
|
|
|
|
text = re.sub(r'(\…{2})([^"’”」』])', r"\1\n\2", text) # 中文省略号
|
|
|
|
|
|
text = re.sub(r'([;；!?。！？\?]["’”」』]{0,2})([^;；!?，。！？\?])', r'\1\n\2', text)
|
|
|
|
|
|
# 如果双引号前有终止符，那么双引号才是句子的终点，把分句符\n放到双引号后，注意前面的几句都小心保留了双引号
|
|
|
|
|
|
text = text.rstrip() # 段尾如果有多余的\n就去掉它
|
|
|
|
|
|
# 很多规则中会考虑分号;，但是这里我把它忽略不计，破折号、英文双引号等同样忽略，需要的再做些简单调整即可。
|
|
|
|
|
|
ls = [i for i in text.split("\n") if i]
|
|
|
|
|
|
for ele in ls:
|
|
|
|
|
|
if len(ele) > self.sentence_size:
|
|
|
|
|
|
ele1 = re.sub(r'([,，.]["’”」』]{0,2})([^,，.])', r'\1\n\2', ele)
|
|
|
|
|
|
ele1_ls = ele1.split("\n")
|
|
|
|
|
|
for ele_ele1 in ele1_ls:
|
|
|
|
|
|
if len(ele_ele1) > self.sentence_size:
|
|
|
|
|
|
ele_ele2 = re.sub(r'([\n]{1,}| {2,}["’”」』]{0,2})([^\s])', r'\1\n\2', ele_ele1)
|
|
|
|
|
|
ele2_ls = ele_ele2.split("\n")
|
|
|
|
|
|
for ele_ele2 in ele2_ls:
|
|
|
|
|
|
if len(ele_ele2) > self.sentence_size:
|
|
|
|
|
|
ele_ele3 = re.sub('( ["’”」』]{0,2})([^ ])', r'\1\n\2', ele_ele2)
|
|
|
|
|
|
ele2_id = ele2_ls.index(ele_ele2)
|
|
|
|
|
|
ele2_ls = ele2_ls[:ele2_id] + [i for i in ele_ele3.split("\n") if i] + ele2_ls[
|
|
|
|
|
|
ele2_id + 1:]
|
|
|
|
|
|
ele_id = ele1_ls.index(ele_ele1)
|
|
|
|
|
|
ele1_ls = ele1_ls[:ele_id] + [i for i in ele2_ls if i] + ele1_ls[ele_id + 1:]
|
|
|
|
|
|
|
|
|
|
|
|
id = ls.index(ele)
|
|
|
|
|
|
ls = ls[:id] + [i for i in ele1_ls if i] + ls[id + 1:]
|
|
|
|
|
|
return ls
|
|
|
|
|
|
|
|
|
|
|
|
class UnstructuredPaddlePDFLoader(UnstructuredFileLoader):
|
|
|
|
|
|
"""Loader that uses unstructured to load image files, such as PNGs and JPGs."""
|
|
|
|
|
|
|
|
|
|
|
|
def _get_elements(self) -> List:
|
|
|
|
|
|
def pdf_ocr_txt(filepath, dir_path="tmp_files"):
|
|
|
|
|
|
full_dir_path = os.path.join(os.path.dirname(filepath), dir_path)
|
|
|
|
|
|
if not os.path.exists(full_dir_path):
|
|
|
|
|
|
os.makedirs(full_dir_path)
|
|
|
|
|
|
ocr = PaddleOCR(use_angle_cls=True, lang="ch", use_gpu=False, show_log=False)
|
|
|
|
|
|
doc = fitz.open(filepath)
|
|
|
|
|
|
txt_file_path = os.path.join(full_dir_path, f"{os.path.split(filepath)[-1]}.txt")
|
|
|
|
|
|
img_name = os.path.join(full_dir_path, 'tmp.png')
|
|
|
|
|
|
with open(txt_file_path, 'w', encoding='utf-8') as fout:
|
|
|
|
|
|
for i in range(doc.page_count):
|
|
|
|
|
|
page = doc[i]
|
|
|
|
|
|
text = page.get_text("")
|
|
|
|
|
|
fout.write(text)
|
|
|
|
|
|
fout.write("\n")
|
|
|
|
|
|
|
|
|
|
|
|
img_list = page.get_images()
|
|
|
|
|
|
for img in img_list:
|
|
|
|
|
|
pix = fitz.Pixmap(doc, img[0])
|
|
|
|
|
|
if pix.n - pix.alpha >= 4:
|
|
|
|
|
|
pix = fitz.Pixmap(fitz.csRGB, pix)
|
|
|
|
|
|
pix.save(img_name)
|
|
|
|
|
|
|
|
|
|
|
|
result = ocr.ocr(img_name)
|
|
|
|
|
|
ocr_result = [i[1][0] for line in result for i in line]
|
|
|
|
|
|
fout.write("\n".join(ocr_result))
|
|
|
|
|
|
if os.path.exists(img_name):
|
|
|
|
|
|
os.remove(img_name)
|
|
|
|
|
|
return txt_file_path
|
|
|
|
|
|
|
|
|
|
|
|
txt_file_path = pdf_ocr_txt(self.file_path)
|
|
|
|
|
|
from unstructured.partition.text import partition_text
|
|
|
|
|
|
return partition_text(filename=txt_file_path, **self.unstructured_kwargs)
|
|
|
|
|
|
|
|
|
|
|
|
def load_caidan(filepath, sentence_size=args.chunk_size):
|
|
|
|
|
|
|
|
|
|
|
|
if filepath.lower().endswith(".pdf"):
|
|
|
|
|
|
# 暂且将paddle相关的loader改为动态加载，可以在不上传pdf/image知识文件的前提下使用protobuf=4.x
|
|
|
|
|
|
loader = UnstructuredPaddlePDFLoader(filepath)
|
|
|
|
|
|
textsplitter = ChineseTextSplitter(pdf=True, sentence_size=sentence_size)
|
|
|
|
|
|
docs = loader.load_and_split(textsplitter)
|
|
|
|
|
|
elif filepath.lower().endswith(".jpg") or filepath.lower().endswith(".png"):
|
|
|
|
|
|
# 暂且将paddle相关的loader改为动态加载，可以在不上传pdf/image知识文件的前提下使用protobuf=4.x
|
|
|
|
|
|
loader = UnstructuredPaddleImageLoader(filepath, mode="elements")
|
|
|
|
|
|
textsplitter = ChineseTextSplitter(pdf=False, sentence_size=sentence_size)
|
|
|
|
|
|
docs = loader.load_and_split(text_splitter=textsplitter)
|
|
|
|
|
|
else:
|
|
|
|
|
|
loader = UnstructuredFileLoader(filepath, mode="elements")
|
|
|
|
|
|
textsplitter = ChineseTextSplitter(pdf=False, sentence_size=sentence_size)
|
|
|
|
|
|
docs = loader.load_and_split(text_splitter=textsplitter)
|
|
|
|
|
|
return docs
|
|
|
|
|
|
|
|
|
|
|
|
def tupian_shibie(mode1,filepath,chatbot):
|
|
|
|
|
|
from aip import AipImageClassify
|
|
|
|
|
|
""" 可选参数 """
|
|
|
|
|
|
options = {}
|
|
|
|
|
|
options["top_num"] = args.top_num
|
|
|
|
|
|
options["filter_threshold"] = str(args.filter_threshold)
|
|
|
|
|
|
options["baike_num"] = args.baike_num
|
|
|
|
|
|
|
|
|
|
|
|
image = get_file_content(filepath)
|
|
|
|
|
|
if mode1 == '菜品识别':
|
|
|
|
|
|
APP_ID = args.APP_ID_c
|
|
|
|
|
|
API_KEY = args.API_KEY_c
|
|
|
|
|
|
SECRET_KEY = args.SECRET_KEY_c
|
|
|
|
|
|
AipImageClassify = AipImageClassify(APP_ID, API_KEY, SECRET_KEY)
|
|
|
|
|
|
result = AipImageClassify.dishDetect(image, options)
|
|
|
|
|
|
result = result["result"]
|
|
|
|
|
|
num = len(result)
|
|
|
|
|
|
name_list = []
|
|
|
|
|
|
probability = []
|
|
|
|
|
|
baike_url = []
|
|
|
|
|
|
description = []
|
|
|
|
|
|
for i in result:
|
|
|
|
|
|
name_list.append(i['name'])
|
|
|
|
|
|
probability.append(i['probability'])
|
|
|
|
|
|
try:
|
|
|
|
|
|
baike_url.append(i['baike_info']['baike_url'])
|
|
|
|
|
|
description.append(i['baike_info']['description'])
|
|
|
|
|
|
except:
|
|
|
|
|
|
baike_url.append('没有百科链接')
|
|
|
|
|
|
description.append('没有百科信息')
|
|
|
|
|
|
inf = '根据预测有以下几种可能:'
|
|
|
|
|
|
for i in range(num):
|
|
|
|
|
|
one = f"""
|
|
|
|
|
|
{i + 1}: {name_list[i]}
|
|
|
|
|
|
可能性为：{probability[i]}
|
|
|
|
|
|
百科链接：{baike_url[i]}
|
|
|
|
|
|
相关描述：{description[i]}
|
|
|
|
|
|
"""
|
|
|
|
|
|
inf += one + "\n"
|
|
|
|
|
|
|
|
|
|
|
|
else:
|
|
|
|
|
|
APP_ID = args.APP_ID_g
|
|
|
|
|
|
API_KEY = args.API_KEY_g
|
|
|
|
|
|
SECRET_KEY = args.SECRET_KEY_g
|
|
|
|
|
|
AipImageClassify = AipImageClassify(APP_ID, API_KEY, SECRET_KEY)
|
|
|
|
|
|
result = AipImageClassify.ingredient(image, options)
|
|
|
|
|
|
|
|
|
|
|
|
result = result["result"]
|
|
|
|
|
|
num = len(result)
|
|
|
|
|
|
name_list = []
|
|
|
|
|
|
score = []
|
|
|
|
|
|
for i in result:
|
|
|
|
|
|
name_list.append(i['name'])
|
|
|
|
|
|
score.append(i['score'])
|
|
|
|
|
|
|
|
|
|
|
|
inf = '根据预测有以下几种可能:'
|
|
|
|
|
|
for i in range(num):
|
|
|
|
|
|
one = f"""
|
|
|
|
|
|
{i + 1}: {name_list[i]}
|
|
|
|
|
|
可能性为：{score[i]}
|
|
|
|
|
|
"""
|
|
|
|
|
|
inf += one + "\n"
|
|
|
|
|
|
|
|
|
|
|
|
b = [None, inf]
|
|
|
|
|
|
chatbot.append(b)
|
|
|
|
|
|
|
|
|
|
|
|
return chatbot
|
|
|
|
|
|
def get_file_content(filePath):
|
|
|
|
|
|
with open(filePath, 'rb') as fp:
|
|
|
|
|
|
return fp.read()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
