Spaces:
Runtime error
Runtime error
| 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() | |