#!python # -*- coding: utf-8 -*- # @author: Kun import gradio as gr import random from sentence_transformers import SentenceTransformer from human_simulator import Human from prompts.service_init import get_init_prompt from utils import get_init, parse_instructions from global_config import lang_opt, llm_model_opt if "openai" == llm_model_opt: from recurrentgpt import RecurrentGPT as AIWriter llm_model = None llm_tokenizer = None elif "vicuna" == llm_model_opt: from recurrent_llm import RecurrentLLM as AIWriter from models.vicuna_bin import load_model llm_tokenizer, llm_model = load_model() elif "chatglm" == llm_model_opt: from recurrent_llm import RecurrentLLM as AIWriter from models.chatglm_hf import load_model llm_tokenizer, llm_model = load_model() elif "baichuan" == llm_model_opt: from recurrent_llm import RecurrentLLM as AIWriter from models.baichuan_hf import load_model llm_tokenizer, llm_model = load_model() elif "aquila" == llm_model_opt: from recurrent_llm import RecurrentLLM as AIWriter from models.aquila_fa import load_model # from models.aquila_hf import load_model llm_tokenizer, llm_model = load_model() elif "falcon" == llm_model_opt: from recurrent_llm import RecurrentLLM from models.falcon_hf import load_model llm_tokenizer, llm_model = load_model() else: raise Exception("not supported llm model name: {}".format(llm_model_opt)) # from urllib.parse import quote_plus # from pymongo import MongoClient # uri = "mongodb://%s:%s@%s" % (quote_plus("xxx"), # quote_plus("xxx"), "localhost") # client = MongoClient(uri, maxPoolSize=None) # db = client.recurrentGPT_db # log = db.log _CACHE = {} # Build the semantic search model embedder = SentenceTransformer('multi-qa-mpnet-base-cos-v1') def init_prompt(novel_type, description): if description == "": description = "" else: description = " about " + description return get_init_prompt(lang_opt, novel_type, description) def init(novel_type, description, request: gr.Request): if novel_type == "": novel_type = "Science Fiction" if "en" == lang_opt else "科幻故事" global _CACHE cookie = request.headers['cookie'] cookie = cookie.split('; _gat_gtag')[0] # prepare first init init_paragraphs = get_init(text=init_prompt( novel_type, description), model=llm_model, tokenizer=llm_tokenizer) # print(init_paragraphs) start_input_to_human = { 'output_paragraph': init_paragraphs['Paragraph 3'], 'input_paragraph': '\n\n'.join([init_paragraphs['Paragraph 1'], init_paragraphs['Paragraph 2'], init_paragraphs['Paragraph 3']]), 'output_memory': init_paragraphs['Summary'], "output_instruction": [init_paragraphs['Instruction 1'], init_paragraphs['Instruction 2'], init_paragraphs['Instruction 3']] } _CACHE[cookie] = {"start_input_to_human": start_input_to_human, "init_paragraphs": init_paragraphs} written_paras = f"""Title: {init_paragraphs['name']} Outline: {init_paragraphs['Outline']} Paragraphs: {start_input_to_human['input_paragraph']}""" if "en" == lang_opt else f"""标题: {init_paragraphs['name']} 梗概: {init_paragraphs['Outline']} 段落: {start_input_to_human['input_paragraph']}""" long_memory = parse_instructions( [init_paragraphs['Paragraph 1'], init_paragraphs['Paragraph 2'], init_paragraphs['Paragraph 3']]) # short memory, long memory, current written paragraphs, 3 next instructions return start_input_to_human['output_memory'], long_memory, written_paras, init_paragraphs['Instruction 1'], init_paragraphs['Instruction 2'], init_paragraphs['Instruction 3'] def step(short_memory, long_memory, instruction1, instruction2, instruction3, current_paras, request: gr.Request, ): if current_paras == "": return "", "", "", "", "", "" global _CACHE # print(list(_CACHE.keys())) # print(request.headers.get('cookie')) cookie = request.headers['cookie'] cookie = cookie.split('; _gat_gtag')[0] cache = _CACHE[cookie] if "writer" not in cache: start_input_to_human = cache["start_input_to_human"] start_input_to_human['output_instruction'] = [ instruction1, instruction2, instruction3] init_paragraphs = cache["init_paragraphs"] human = Human(input=start_input_to_human, memory=None, embedder=embedder, model=llm_model, tokenizer=llm_tokenizer) human.step() start_short_memory = init_paragraphs['Summary'] writer_start_input = human.output # Init writerGPT writer = AIWriter(input=writer_start_input, short_memory=start_short_memory, long_memory=[ init_paragraphs['Paragraph 1'], init_paragraphs['Paragraph 2'], init_paragraphs['Paragraph 3']], memory_index=None, embedder=embedder, model=llm_model, tokenizer=llm_tokenizer) cache["writer"] = writer cache["human"] = human writer.step() else: human = cache["human"] writer = cache["writer"] output = writer.output output['output_memory'] = short_memory # randomly select one instruction out of three instruction_index = random.randint(0, 2) output['output_instruction'] = [instruction1, instruction2, instruction3][instruction_index] human.input = output human.step() writer.input = human.output writer.step() long_memory = [[v] for v in writer.long_memory] # short memory, long memory, current written paragraphs, 3 next instructions return writer.output['output_memory'], long_memory, current_paras + '\n\n' + writer.output['input_paragraph'], human.output['output_instruction'], *writer.output['output_instruction'] def controled_step(short_memory, long_memory, selected_instruction, current_paras, request: gr.Request, ): if current_paras == "": return "", "", "", "", "", "" global _CACHE # print(list(_CACHE.keys())) # print(request.headers.get('cookie')) cookie = request.headers['cookie'] cookie = cookie.split('; _gat_gtag')[0] cache = _CACHE[cookie] if "writer" not in cache: start_input_to_human = cache["start_input_to_human"] start_input_to_human['output_instruction'] = selected_instruction init_paragraphs = cache["init_paragraphs"] human = Human(input=start_input_to_human, memory=None, embedder=embedder, model=llm_model, tokenizer=llm_tokenizer) human.step() start_short_memory = init_paragraphs['Summary'] writer_start_input = human.output # Init writerGPT writer = AIWriter(input=writer_start_input, short_memory=start_short_memory, long_memory=[ init_paragraphs['Paragraph 1'], init_paragraphs['Paragraph 2'], init_paragraphs['Paragraph 3']], memory_index=None, embedder=embedder, model=llm_model, tokenizer=llm_tokenizer) cache["writer"] = writer cache["human"] = human writer.step() else: human = cache["human"] writer = cache["writer"] output = writer.output output['output_memory'] = short_memory output['output_instruction'] = selected_instruction human.input = output human.step() writer.input = human.output writer.step() # short memory, long memory, current written paragraphs, 3 next instructions return writer.output['output_memory'], parse_instructions(writer.long_memory), current_paras + '\n\n' + writer.output['input_paragraph'], *writer.output['output_instruction'] # SelectData is a subclass of EventData def on_select(instruction1, instruction2, instruction3, evt: gr.SelectData): selected_plan = int(evt.value.replace("Instruction ", "") ) if "en" == lang_opt else int(evt.value.replace("指令 ", "")) selected_plan = [instruction1, instruction2, instruction3][selected_plan-1] return selected_plan def reload_model(choice): pass with gr.Blocks(title="小说GPT", css="footer {visibility: hidden}", theme="default") as demo: if "en" == lang_opt: gr.Markdown( """ # Recurrent-LLM Interactive Generation of (Arbitrarily) Long Texts with Human-in-the-Loop """) elif lang_opt in ["zh1", "zh2"]: gr.Markdown( """ # Recurrent-LLM 可以根据题目和简介自动续写文章 也可以手动选择剧情走向进行续写 """) with gr.Tab("Auto-Generation"): with gr.Row(): with gr.Column(): with gr.Row(): with gr.Column(scale=1, min_width=200): novel_type = gr.Textbox( label="Novel Type", placeholder="e.g. science fiction") if "en" == lang_opt else gr.Textbox( label="请输入文本", placeholder="可以自己填写或者从EXamples中选择一个填入") with gr.Column(scale=2, min_width=400): description = gr.Textbox( label="Description") if "en" == lang_opt else gr.Textbox(label="剧情简介（非必选项）") btn_init = gr.Button( "Init Novel Generation", variant="primary") if "en" == lang_opt else gr.Button( "点击开始运行", variant="primary") if "en" == lang_opt: gr.Examples(["Science Fiction", "Romance", "Mystery", "Fantasy", "Historical", "Horror", "Thriller", "Western", "Young Adult", ], inputs=[novel_type]) elif lang_opt in ["zh1", "zh2"]: gr.Examples(["科幻故事", "青春伤痛文学", "爱到死去活来", "搞笑", "幽默", "鬼故事", "喜剧", "童话", "魔法世界", ], inputs=[novel_type]) written_paras = gr.Textbox( label="Written Paragraphs (editable)", max_lines=21, lines=21) if "en" == lang_opt else gr.Textbox( label="文章内容", max_lines=21, lines=21) with gr.Column(): if "en" == lang_opt: gr.Markdown("### Memory Module\n") elif lang_opt in ["zh1", "zh2"]: gr.Markdown("### 剧情模型\n") short_memory = gr.Textbox( label="Short-Term Memory (editable)", max_lines=3, lines=3) if "en" == lang_opt else gr.Textbox( label="短期记忆 (可编辑)", max_lines=3, lines=3) long_memory = gr.Textbox( label="Long-Term Memory (editable)", max_lines=6, lines=6) if "en" == lang_opt else gr.Textbox( label="长期记忆 (可编辑)", max_lines=6, lines=6) if "en" == lang_opt: gr.Markdown("### Instruction Module\n") elif lang_opt in ["zh1", "zh2"]: gr.Markdown("### 选项模型\n") with gr.Row(): instruction1 = gr.Textbox( label="Instruction 1 (editable)", max_lines=4, lines=4) if "en" == lang_opt else gr.Textbox( label="指令1(可编辑)", max_lines=4, lines=4) instruction2 = gr.Textbox( label="Instruction 2 (editable)", max_lines=4, lines=4) if "en" == lang_opt else gr.Textbox( label="指令2(可编辑)", max_lines=4, lines=4) instruction3 = gr.Textbox( label="Instruction 3 (editable)", max_lines=4, lines=4) if "en" == lang_opt else gr.Textbox( label="指令3(可编辑)", max_lines=4, lines=4) selected_plan = gr.Textbox( label="Revised Instruction (from last step)", max_lines=2, lines=2) if "en" == lang_opt else gr.Textbox( label="选项说明 (来自上一步)", max_lines=2, lines=2) btn_step = gr.Button("Next Step", variant="primary") if "en" == lang_opt else gr.Button( "下一步", variant="primary") btn_init.click(init, inputs=[novel_type, description], outputs=[ short_memory, long_memory, written_paras, instruction1, instruction2, instruction3]) btn_step.click(step, inputs=[short_memory, long_memory, instruction1, instruction2, instruction3, written_paras], outputs=[ short_memory, long_memory, written_paras, selected_plan, instruction1, instruction2, instruction3]) with gr.Tab("Human-in-the-Loop"): with gr.Row(): with gr.Column(): with gr.Row(): with gr.Column(scale=1, min_width=200): novel_type = gr.Textbox( label="Novel Type", placeholder="e.g. science fiction") if "en" == lang_opt else gr.Textbox( label="请输入文本", placeholder="可以自己填写或者从EXamples中选择一个填入") with gr.Column(scale=2, min_width=400): description = gr.Textbox( label="Description") if "en" == lang_opt else gr.Textbox(label="剧情简介（非必选项）") btn_init = gr.Button( "Init Novel Generation", variant="primary") if "en" == lang_opt else gr.Button( "点击开始运行", variant="primary") if "en" == lang_opt: gr.Examples(["Science Fiction", "Romance", "Mystery", "Fantasy", "Historical", "Horror", "Thriller", "Western", "Young Adult", ], inputs=[novel_type]) elif lang_opt in ["zh1", "zh2"]: gr.Examples(["科幻小说", "爱情小说", "推理小说", "奇幻小说", "玄幻小说", "恐怖", "悬疑", "惊悚", "武侠小说", ], inputs=[novel_type]) written_paras = gr.Textbox( label="Written Paragraphs (editable)", max_lines=23, lines=23) if "en" == lang_opt else gr.Textbox( label="文章内容 (可编辑)", max_lines=23, lines=23) with gr.Column(): if "en" == lang_opt: gr.Markdown("### Memory Module\n") elif lang_opt in ["zh1", "zh2"]: gr.Markdown("### 剧情模型\n") short_memory = gr.Textbox( label="Short-Term Memory (editable)", max_lines=3, lines=3) if "en" == lang_opt else gr.Textbox( label="短期记忆 (可编辑)", max_lines=3, lines=3) long_memory = gr.Textbox( label="Long-Term Memory (editable)", max_lines=6, lines=6) if "en" == lang_opt else gr.Textbox( label="长期记忆 (可编辑)", max_lines=6, lines=6) if "en" == lang_opt: gr.Markdown("### Instruction Module\n") elif lang_opt in ["zh1", "zh2"]: gr.Markdown("### 选项模型\n") with gr.Row(): instruction1 = gr.Textbox( label="Instruction 1", max_lines=3, lines=3, interactive=False) if "en" == lang_opt else gr.Textbox( label="指令1", max_lines=3, lines=3, interactive=False) instruction2 = gr.Textbox( label="Instruction 2", max_lines=3, lines=3, interactive=False) if "en" == lang_opt else gr.Textbox( label="指令2", max_lines=3, lines=3, interactive=False) instruction3 = gr.Textbox( label="Instruction 3", max_lines=3, lines=3, interactive=False) if "en" == lang_opt else gr.Textbox( label="指令3", max_lines=3, lines=3, interactive=False) with gr.Row(): with gr.Column(scale=1, min_width=100): selected_plan = gr.Radio( ["Instruction 1", "Instruction 2", "Instruction 3"], label="Instruction Selection",) if "en" == lang_opt else gr.Radio(["指令 1", "指令 2", "指令 3"], label="指令 选择",) with gr.Column(scale=3, min_width=300): selected_instruction = gr.Textbox( label="Selected Instruction (editable)", max_lines=5, lines=5) if "en" == lang_opt else gr.Textbox( label="在上一步骤中被选择的 (可编辑)", max_lines=5, lines=5) btn_step = gr.Button("Next Step", variant="primary") if "en" == lang_opt else gr.Button( "下一步", variant="primary") btn_init.click(init, inputs=[novel_type, description], outputs=[ short_memory, long_memory, written_paras, instruction1, instruction2, instruction3]) btn_step.click(controled_step, inputs=[short_memory, long_memory, selected_instruction, written_paras], outputs=[ short_memory, long_memory, written_paras, instruction1, instruction2, instruction3]) selected_plan.select(on_select, inputs=[ instruction1, instruction2, instruction3], outputs=[selected_instruction]) with gr.Tab("Model-Config"): model_opt_radio = gr.Radio(["OpenAI", "ChatGLM-6B", "Vicuna-7B"], value="OpenAI", label="model", info="select language you preferred. Default is English.", interactive=True ) reload_button = gr.Button("Reload/重新加载") reload_button.click(reload_model, show_progress=True, inputs=[model_opt_radio], outputs=[novel_type]) if __name__ == "__main__": demo.queue().launch(show_error=True, show_api=False)