import requests import json import gradio as gr import pdfplumber import pandas as pd import time from cnocr import CnOcr import numpy as np import openai from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader, Prompt from transformers import pipeline, BarkModel, BarkProcessor import opencc import scipy import torch import hashlib converter = opencc.OpenCC('t2s') # 创建一个OpenCC实例，指定繁体字转为简体字 ocr = CnOcr() # 初始化ocr模型 history_max_len = 500 # 机器人记忆的最大长度 all_max_len = 2000 # 输入的最大长度 asr_model_id = "openai/whisper-tiny" # 更新为你的模型ID device = "cuda:0" if torch.cuda.is_available() else "cpu" asr_pipe = pipeline("automatic-speech-recognition", model=asr_model_id, device=device) bark_model = BarkModel.from_pretrained("suno/bark-small") bark_processor = BarkProcessor.from_pretrained("suno/bark-small") sampling_rate = bark_model.generation_config.sample_rate def get_text_emb(open_ai_key, text): # 文本向量化 openai.api_key = open_ai_key # 设置openai的key response = openai.Embedding.create( input=text, model="text-embedding-ada-002" ) # 调用openai的api return response['data'][0]['embedding'] # 返回向量 def doc_index_self(open_ai_key, doc): # 文档向量化 texts = doc.split('\n') # 按行切分 emb_list = [] # 用于存储向量 for text in texts: # 遍历每一行 emb_list.append(get_text_emb(open_ai_key, text)) # 获取向量 return texts, emb_list, gr.Textbox.update(visible=True), gr.Button.update(visible=True), gr.Markdown.update( value="""操作说明 step 3：建立索引（by self）成功！ 🙋 可以开始对话啦~"""), gr.Chatbot.update(visible=True), 1, gr.Audio.update( visible=True), gr.Radio.update(visible=True) def doc_index_llama(open_ai_key, txt): # 建立索引 # 根据时间戳新建目录，保存txt文件 path = str(time.time()) import os os.mkdir(path) with open(path + '/doc.txt', mode='w', encoding='utf-8') as f: f.write(txt) openai.api_key = open_ai_key # 设置OpenAI API Key documents = SimpleDirectoryReader(path).load_data() # 读取文档 index = GPTVectorStoreIndex.from_documents(documents) # 建立索引 template = ( "你是一个有用的助手，可以使用文章内容准确地回答问题。使用提供的文章来生成你的答案，但避免逐字复制文章。尽可能使用自己的话。准确、有用、简洁、清晰。文章内容如下： \n" "---------------------\n" "{context_str}" "\n---------------------\n" "{query_str}\n" "请你回复用户。\n" ) # 定义模板 qa_template = Prompt(template) # 将模板转换成Prompt对象 query_engine = index.as_query_engine(text_qa_template=qa_template) # 建立查询引擎 return query_engine, gr.Textbox.update(visible=True), gr.Button.update(visible=True), gr.Markdown.update( value="""操作说明 step 3：建立索引（by llama_index）成功！ 🙋 可以开始对话啦~"""), gr.Chatbot.update( visible=True), 0, gr.Audio.update(visible=True), gr.Radio.update(visible=True) def get_response_by_self(open_ai_key, msg, bot, doc_text_list, doc_embeddings): # 获取机器人回复 now_len = len(msg) # 当前输入的长度 his_bg = -1 # 历史记录的起始位置 for i in range(len(bot) - 1, -1, -1): # 从后往前遍历历史记录 if now_len + len(bot[i][0]) + len(bot[i][1]) > history_max_len: # 如果超过了历史记录的最大长度，就不再加入 break now_len += len(bot[i][0]) + len(bot[i][1]) # 更新当前长度 his_bg = i # 更新历史记录的起始位置 history = [] if his_bg == -1 else bot[his_bg:] # 获取历史记录 query_embedding = get_text_emb(open_ai_key, msg) # 获取输入的向量 cos_scores = [] # 用于存储相似度 def cos_sim(a, b): # 计算余弦相似度 return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)) # 返回相似度 for doc_embedding in doc_embeddings: # 遍历文档向量 cos_scores.append(cos_sim(query_embedding, doc_embedding)) # 计算相似度 score_index = [] # 用于存储相似度和索引对应 for i in range(len(cos_scores)): # 遍历相似度 score_index.append((cos_scores[i], i)) # 加入相似度和索引对应 score_index.sort(key=lambda x: x[0], reverse=True) # 按相似度排序 print('score_index:\n', score_index) index_set, sub_doc_list = set(), [] # 用于存储最终的索引和文档 for s_i in score_index: # 遍历相似度和索引对应 doc = doc_text_list[s_i[1]] # 获取文档 if now_len + len(doc) > all_max_len: # 如果超过了最大长度，就不再加入 break index_set.add(s_i[1]) # 加入索引 now_len += len(doc) # 更新当前长度 # 可能段落截断错误，所以把上下段也加入进来 if s_i[1] > 0 and s_i[1] - 1 not in index_set: # 如果上一段没有加入 doc = doc_text_list[s_i[1] - 1] # 获取上一段 if now_len + len(doc) > all_max_len: # 如果超过了最大长度，就不再加入 break index_set.add(s_i[1] - 1) # 加入索引 now_len += len(doc) # 更新当前长度 if s_i[1] + 1 < len(doc_text_list) and s_i[1] + 1 not in index_set: # 如果下一段没有加入 doc = doc_text_list[s_i[1] + 1] # 获取下一段 if now_len + len(doc) > all_max_len: # 如果超过了最大长度，就不再加入 break index_set.add(s_i[1] + 1) # 加入索引 now_len += len(doc) # 更新当前长度 index_list = list(index_set) # 转换成list index_list.sort() # 排序 for i in index_list: # 遍历索引 sub_doc_list.append(doc_text_list[i]) # 加入文档 document = '' if len(sub_doc_list) == 0 else '\n'.join(sub_doc_list) # 拼接文档 messages = [{ "role": "system", "content": "你是一个有用的助手，可以使用文章内容准确地回答问题。使用提供的文章来生成你的答案，但避免逐字复制文章。尽可能使用自己的话。准确、有用、简洁、清晰。" }, {"role": "system", "content": "文章内容：\n" + document}] # 角色人物定义 for his in history: # 遍历历史记录 messages.append({"role": "user", "content": his[0]}) # 加入用户的历史记录 messages.append({"role": "assistant", "content": his[1]}) # 加入机器人的历史记录 messages.append({"role": "user", "content": msg}) # 加入用户的当前输入 openai.api_key = open_ai_key chat_completion = openai.ChatCompletion.create(model="gpt-3.5-turbo", messages=messages) # 获取机器人的回复 res = chat_completion.choices[0].message.content # 获取机器人的回复 bot.append([msg, res]) # 加入历史记录 return bot[max(0, len(bot) - 3):] # 返回最近3轮的历史记录 def get_response_by_llama_index(open_ai_key, msg, bot, query_engine): # 获取机器人回复 openai.api_key = open_ai_key query_str = "历史对话如下：\n" for his in bot: # 遍历历史记录 query_str += "用户：" + his[0] + "\n" # 加入用户的历史记录 query_str += "助手：" + his[1] + "\n" # 加入机器人的历史记录 query_str += "用户：" + msg + "\n" # 加入用户的当前输入 res = query_engine.query(query_str) # 获取回答 print(res) # 显示回答 bot.append([msg, str(res)]) # 加入历史记录 return bot[max(0, len(bot) - 3):] # 返回最近3轮的历史记录 def get_audio_answer(bot): # 获取语音回答 answer = bot[-1][1] inputs = bark_processor( text=[answer], return_tensors="pt", ) speech_values = bark_model.generate(**inputs, do_sample=True) au_dir = hashlib.md5(answer.encode('utf-8')).hexdigest() + '.wav' # 获取md5 scipy.io.wavfile.write(au_dir, rate=sampling_rate, data=speech_values.cpu().numpy().squeeze()) return gr.Audio().update(au_dir, autoplay=True) def get_response(open_ai_key, msg, bot, doc_text_list, doc_embeddings, query_engine, index_type): # 获取机器人回复 if index_type == 1: # 如果是使用自己的索引 bot = get_response_by_self(open_ai_key, msg, bot, doc_text_list, doc_embeddings) else: # 如果是使用llama_index索引 bot = get_response_by_llama_index(open_ai_key, msg, bot, query_engine) return bot def up_file(files): # 上传文件 doc_text_list = [] # 用于存储文档 for idx, file in enumerate(files): # 遍历文件 print(file.name) with pdfplumber.open(file.name) as pdf: # 打开pdf for i in range(len(pdf.pages)): # 遍历pdf的每一页 # 读取PDF文档第i+1页 page = pdf.pages[i] res_list = page.extract_text().split('\n')[:-1] # 提取文本 for j in range(len(page.images)): # 遍历图片 # 获取图片的二进制流 img = page.images[j] file_name = '{}-{}-{}.png'.format(str(time.time()), str(i), str(j)) # 生成文件名 with open(file_name, mode='wb') as f: # 保存图片 f.write(img['stream'].get_data()) try: res = ocr.ocr(file_name) # 识别图片 except Exception as e: res = [] # 识别失败 if len(res) > 0: # 如果识别成功 res_list.append(' '.join([re['text'] for re in res])) # 加入识别结果 tables = page.extract_tables() # 提取表格 for table in tables: # 遍历表格 # 第一列当成表头： df = pd.DataFrame(table[1:], columns=table[0]) try: records = json.loads(df.to_json(orient="records", force_ascii=False)) # 转换成json for rec in records: # 遍历json res_list.append(json.dumps(rec, ensure_ascii=False)) # 加入json except Exception as e: res_list.append(str(df)) # 如果转换识别，直接把表格转为str doc_text_list += res_list # 加入文档 doc_text_list = [str(text).strip() for text in doc_text_list if len(str(text).strip()) > 0] # 去除空格 print(doc_text_list) return gr.Textbox.update(value='\n'.join(doc_text_list), visible=True), gr.Button.update( visible=True), gr.Button.update( visible=True), gr.Markdown.update( value="操作说明 step 2：确认PDF解析结果（可修正），点击“建立索引”，随后进行对话") def transcribe_speech_by_self(filepath): output = asr_pipe( filepath, max_new_tokens=256, generate_kwargs={ "task": "transcribe", "language": "chinese", }, chunk_length_s=30, batch_size=8, ) # 识别语音 simplified_text = converter.convert(output["text"]) # 转换为简体字 return simplified_text def transcribe_speech_by_openai(openai_key, filepath): openai.api_key = openai_key # 设置OpenAI API Key audio_file = open(filepath, "rb") transcript = openai.Audio.transcribe("whisper-1", audio_file) print(transcript) return transcript['text'] def transcribe_speech(openai_key, filepath, a_type): if a_type == 'self': return transcribe_speech_by_self(filepath) else: return transcribe_speech_by_openai(openai_key, filepath) with gr.Blocks() as demo: with gr.Row(): with gr.Column(): open_ai_key = gr.Textbox(label='OpenAI API Key', placeholder='输入你的OpenAI API Key') # 你的OpenAI API Key file = gr.File(file_types=['.pdf'], label='点击上传PDF，进行解析(支持多文档、表格、OCR)', file_count='multiple') # 支持多文档、表格、OCR txt = gr.Textbox(label='PDF解析结果', visible=False) # PDF解析结果 with gr.Row(): index_llama_bu = gr.Button(value='建立索引（by llama_index）', visible=False) # 建立索引（by llama_index） index_self_bu = gr.Button(value='建立索引（by self）', visible=False) # 建立索引（by self） doc_text_state = gr.State([]) # 存储PDF解析结果 doc_emb_state = gr.State([]) # 存储PDF解析结果的embedding query_engine = gr.State([]) # 存储查询引擎 index_type = gr.State([]) # 存储索引类型 with gr.Column(): md = gr.Markdown("""操作说明 step 1：点击左侧区域，上传PDF，进行解析""") # 操作说明 chat_bot = gr.Chatbot(visible=False) # 聊天机器人 audio_answer = gr.Audio() # 语音回答 with gr.Row(): asr_type = gr.Radio(value='self', choices=['self', 'openai'], label='语音识别方式', visible=False) # 语音识别方式 audio_inputs = gr.Audio(source="microphone", type="filepath", label="点击录音输入", visible=False) # 录音输入 msg_txt = gr.Textbox(label='消息框', placeholder='输入消息', visible=False) # 消息框 chat_bu = gr.Button(value='发送', visible=False) # 发送按钮 file.change(up_file, [file], [txt, index_self_bu, index_llama_bu, md]) # 上传文件 index_self_bu.click(doc_index_self, [open_ai_key, txt], [doc_text_state, doc_emb_state, msg_txt, chat_bu, md, chat_bot, index_type, audio_inputs, asr_type]) # 提交解析结果 index_llama_bu.click(doc_index_llama, [open_ai_key, txt], [query_engine, msg_txt, chat_bu, md, chat_bot, index_type, audio_inputs, asr_type]) # 提交解析结果 audio_inputs.change(transcribe_speech, [open_ai_key, audio_inputs, asr_type], [msg_txt]) # 录音输入 chat_bu.click(get_response, [open_ai_key, msg_txt, chat_bot, doc_text_state, doc_emb_state, query_engine, index_type], [chat_bot])# .then(get_audio_answer, [chat_bot], [audio_answer]) # 发送消息 if __name__ == "__main__": demo.queue(concurrency_count=4).launch()