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