Upload 33 files

.gitignore

@@ -0,0 +1,160 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 MK
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

app.py

@@ -0,0 +1,391 @@
+#!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="RecurrentGPT", 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.Box():
+                    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])
+                else:
+                    raise Exception(f"not supported language: {lang_opt}")
+
+                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():
+                with gr.Box():
+                    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)
+                    # long_memory = gr.Dataframe(
+                    #     # label="Long-Term Memory (editable)",
+                    #     headers=["Long-Term Memory (editable)"],
+                    #     datatype=["str"],
+                    #     row_count=3,
+                    #     max_rows=3,
+                    #     col_count=(1, "fixed"),
+                    #     type="array",
+                    # )
+                with gr.Box():
+                    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.Box():
+                    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():
+                with gr.Box():
+                    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)
+                with gr.Box():
+                    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="指令 选择",)
+                            #  info="Select the instruction you want to revise and use for the next step generation.")
+                        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])
+
+    demo.queue(concurrency_count=1)
+
+if __name__ == "__main__":
+    demo.launch(show_error=True,show_api=False)

common/__init__.py

@@ -0,0 +1,19 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+
+import torch
+
+def torch_gc():
+    if torch.cuda.is_available():
+        # with torch.cuda.device(DEVICE):
+        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 产生的内存占用。")

global_config.py

@@ -0,0 +1,20 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+####################################################
+
+# lang_opt = "zh1"
+lang_opt = "zh2"
+# lang_opt = "en"
+
+####################################################
+
+# llm_model_opt = "openai"
+# llm_model_opt = "vicuna"
+llm_model_opt = "chatglm"
+# llm_model_opt = "baichuan"
+# llm_model_opt = "aquila"
+# llm_model_opt = "falcon"
+
+####################################################

human_simulator.py

@@ -0,0 +1,137 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+
+from utils import get_content_between_a_b, parse_instructions
+from prompts.human_simulator import get_input_text
+from global_config import lang_opt, llm_model_opt
+
+if "openai" == llm_model_opt:
+    from utils.openai_util import get_api_response
+elif "vicuna" == llm_model_opt:
+    from utils.vicuna_util import get_api_response
+elif "chatglm" == llm_model_opt:
+    from utils.chatglm_util import get_api_response
+elif "baichuan" == llm_model_opt:
+    from utils.baichuan_util import get_api_response
+elif "aquila" == llm_model_opt:
+    from utils.aquila_util import get_api_response
+elif "falcon" == llm_model_opt:
+    from utils.falcon_util import get_api_response
+else:
+    raise Exception("not supported llm model name: {}".format(llm_model_opt))
+
+
+class Human:
+
+    def __init__(self, input, memory, embedder, model, tokenizer):
+        self.input = input
+        if memory:
+            self.memory = memory
+        else:
+            self.memory = self.input['output_memory']
+        self.embedder = embedder
+        self.model = model
+        self.tokenizer = tokenizer
+        self.output = {}
+
+    def prepare_input(self):
+        previous_paragraph = self.input["input_paragraph"]
+        writer_new_paragraph = self.input["output_paragraph"]
+        memory = self.input["output_memory"]
+        user_edited_plan = self.input["output_instruction"]
+
+        input_text = get_input_text(
+            lang_opt, previous_paragraph, memory, writer_new_paragraph, user_edited_plan)
+
+        return input_text
+
+    def parse_plan(self, response):
+        plan = get_content_between_a_b('Selected Plan:', 'Reason', response)
+        return plan
+
+    def select_plan(self, response_file): # TODO ???
+
+        previous_paragraph = self.input["input_paragraph"]
+        writer_new_paragraph = self.input["output_paragraph"]
+        memory = self.input["output_memory"]
+        previous_plans = self.input["output_instruction"]
+        prompt = f"""
+    Now imagine you are a helpful assistant that help a novelist with decision making. You will be given a previously written paragraph and a paragraph written by a ChatGPT writing assistant, a summary of the main storyline maintained by the ChatGPT assistant, and 3 different possible plans of what to write next.
+    I need you to:
+    Select the most interesting and suitable plan proposed by the ChatGPT assistant.
+
+    Previously written paragraph:  
+    {previous_paragraph}
+
+    The summary of the main storyline maintained by your ChatGPT assistant:
+    {memory}
+
+    The new paragraph written by your ChatGPT assistant:
+    {writer_new_paragraph}
+
+    Three plans of what to write next proposed by your ChatGPT assistant:
+    {parse_instructions(previous_plans)}
+
+    Now start choosing, organize your output by strictly following the output format as below:
+      
+    Selected Plan: 
+    <copy the selected plan here>
+
+    Reason:
+    <Explain why you choose the plan>
+    """
+        print(prompt+'\n'+'\n')
+
+        response = get_api_response(self.model, self.tokenizer, prompt)
+
+        plan = self.parse_plan(response)
+        while plan == None:
+            response = get_api_response(self.model, self.tokenizer, prompt)
+            plan = self.parse_plan(response)
+
+        if response_file:
+            with open(response_file, 'a', encoding='utf-8') as f:
+                f.write(f"Selected plan here:\n{response}\n\n")
+
+        return plan
+
+    def parse_output(self, text):
+        try:
+            if text.splitlines()[0].startswith('Extended Paragraph'):
+                new_paragraph = get_content_between_a_b(
+                    'Extended Paragraph:', 'Selected Plan', text)
+            else:
+                new_paragraph = text.splitlines()[0]
+
+            lines = text.splitlines()
+            if lines[-1] != '\n' and lines[-1].startswith('Revised Plan:'):
+                revised_plan = lines[-1][len("Revised Plan:"):]
+            elif lines[-1] != '\n':
+                revised_plan = lines[-1]
+
+            output = {
+                "output_paragraph": new_paragraph,
+                # "selected_plan": selected_plan,
+                "output_instruction": revised_plan,
+                # "memory":self.input["output_memory"]
+            }
+
+            return output
+        except:
+            return None
+
+    def step(self, response_file=None):
+
+        prompt = self.prepare_input()
+        print(prompt+'\n'+'\n')
+
+        response = get_api_response(self.model, self.tokenizer, prompt)
+        self.output = self.parse_output(response)
+        while self.output == None:
+            response = get_api_response(self.model, self.tokenizer, prompt)
+            self.output = self.parse_output(response)
+        if response_file:
+            with open(response_file, 'a', encoding='utf-8') as f:
+                f.write(f"Human's output here:\n{response}\n\n")

init_prompt.json

@@ -0,0 +1 @@
+{"init_prompt": "\nPlease write a {type} novel about {topic} with about 50 chapters. Follow the format below precisely:\n\n    Begin with the name of the novel.\n    Next, write an outline for the first chapter. The outline should describe the background and the beginning of the novel.\n    Write the first three paragraphs with their indication of the novel based on your outline. Write in a novelistic style and take your time to set the scene.\n    Write a summary that captures the key information of the three paragraphs.\n    Finally, write three different instructions for what to write next, each containing around five sentences. Each instruction should present a possible, interesting continuation of the story.\n    The output format should follow these guidelines:\n    Name: <name of the novel>\n    Outline: <outline for the first chapter>\n    Paragraph 1: <content for paragraph 1>\n    Paragraph 2: <content for paragraph 2>\n    Paragraph 3: <content for paragraph 3>\n    Summary: <content of summary>\n    Instruction 1: <content for instruction 1>\n    Instruction 2: <content for instruction 2>\n    Instruction 3: <content for instruction 3>\n    \n    Make sure to be precise and follow the output format strictly.\n    \n    "}

misc/placeholder

@@ -0,0 +1 @@
+

models/__init__.py

@@ -0,0 +1,4 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+

models/aquila_fa.py

@@ -0,0 +1,38 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+import os
+import torch
+from flagai.auto_model.auto_loader import AutoLoader
+from flagai.model.predictor.predictor import Predictor
+from flagai.model.predictor.aquila import aquila_generate
+from flagai.data.tokenizer import Tokenizer
+import bminf
+
+
+
+max_token: int = 128 # 10000 # 64 
+temperature: float = 0.75
+top_p = 0.9
+
+state_dict = "./checkpoints_in"
+model_name = 'aquilachat-7b'
+
+def load_model():
+    loader = AutoLoader(
+        "lm",
+        model_dir=state_dict,
+        model_name=model_name,
+        use_cache=True,
+        fp16=True)
+    model = loader.get_model()
+    tokenizer = loader.get_tokenizer()
+    cache_dir = os.path.join(state_dict, model_name)
+
+    model.eval()
+
+    with torch.cuda.device(0):
+        model = bminf.wrapper(model, quantization=False, memory_limit=2 << 30)
+        
+    return tokenizer, model

models/aquila_hf.py

@@ -0,0 +1,20 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# trust_remote_code: remote code depends old version transformers
+"""
+File "/root/.cache/huggingface/modules/transformers_modules/qhduan/aquilachat-7b/9d8fcc4f12b6bb6ea0c8a494ba85110f78804739/modeling_aquila.py", line 33, in <module>
+    from transformers.models.llama.configuration_llama import LlamaConfig
+ModuleNotFoundError: No module named 'transformers.models.llama'
+"""
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained('qhduan/aquilachat-7b')
+    model = AutoModelForCausalLM.from_pretrained('qhduan/aquilachat-7b', trust_remote_code=True)
+    model = model.eval().half().cuda()
+
+    return tokenizer, model

models/baichuan_hf.py

@@ -0,0 +1,50 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+
+max_token: int = 10000  # 10000 # 64
+temperature: float = 0.75
+top_p = 0.9
+use_lora = False
+
+
+# def load_model():
+#     model_name_or_path = "baichuan-inc/baichuan-7B"
+#     # model_name_or_path = "~/.cache/huggingface/hub/models--baichuan-inc--baichuan-7B/snapshots/39916f64eb892ccdc1982b0eef845b3b8fd43f6b/"
+#     tokenizer = AutoTokenizer.from_pretrained(
+#         model_name_or_path,
+#         trust_remote_code=True)  
+#     model = AutoModelForCausalLM.from_pretrained(
+#         model_name_or_path,
+#         device_map="auto",
+#         trust_remote_code=True)  
+
+#     # inputs = tokenizer('登鹳雀楼->王之涣\n夜雨寄北->', return_tensors='pt')
+#     # inputs = inputs.to('cuda:0')
+#     # pred = model.generate(**inputs, max_new_tokens=64,repetition_penalty=1.1)
+#     # print(tokenizer.decode(pred.cpu()[0], skip_special_tokens=True))
+
+#     return tokenizer, model
+
+
+def load_model(use_lora=True, LOAD_IN_8BIT=False):
+    """
+    params: 
+    use_lora=True, LOAD_IN_8BIT=False
+    use_lora=False. LOAD_IN_8BIT=True
+    """
+    tokenizer = AutoTokenizer.from_pretrained("baichuan-inc/baichuan-7B",
+                                              trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained("baichuan-inc/baichuan-7B",
+                                                 device_map="auto",
+                                                 trust_remote_code=True,
+                                                 load_in_8bit=LOAD_IN_8BIT, # if not have enough GPU memory, then use 8bit
+                                                 )
+    
+    if use_lora:
+        model = PeftModel.from_pretrained(model, "hiyouga/baichuan-7b-sft")
+
+    return tokenizer, model

models/chatglm_hf.py

@@ -0,0 +1,76 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+import torch
+from transformers import AutoTokenizer, AutoConfig, AutoModel
+
+model_name_or_path = "THUDM/chatglm-6b-int8" 
+max_token: int = 10000
+temperature: float = 0.75
+top_p = 0.9
+use_lora = False
+
+def auto_configure_device_map(num_gpus: int, use_lora: bool):
+    # transformer.word_embeddings 占用1层
+    # transformer.final_layernorm 和 lm_head 占用1层
+    # transformer.layers 占用 28 层
+    # 总共30层分配到num_gpus张卡上
+    num_trans_layers = 28
+    per_gpu_layers = 30 / num_gpus
+
+    # bugfix: PEFT加载lora模型出现的层命名不同
+    # if LLM_LORA_PATH and use_lora:
+    #     layer_prefix = 'base_model.model.transformer'
+    # else:
+    layer_prefix = 'transformer'
+
+    # bugfix: 在linux中调用torch.embedding传入的weight,input不在同一device上,导致RuntimeError
+    # windows下 model.device 会被设置成 transformer.word_embeddings.device
+    # linux下 model.device 会被设置成 lm_head.device
+    # 在调用chat或者stream_chat时,input_ids会被放到model.device上
+    # 如果transformer.word_embeddings.device和model.device不同,则会导致RuntimeError
+    # 因此这里将transformer.word_embeddings,transformer.final_layernorm,lm_head都放到第一张卡上
+    device_map = {f'{layer_prefix}.word_embeddings': 0,
+                  f'{layer_prefix}.final_layernorm': 0, 'lm_head': 0,
+                  f'base_model.model.lm_head': 0, }
+
+    used = 2
+    gpu_target = 0
+    for i in range(num_trans_layers):
+        if used >= per_gpu_layers:
+            gpu_target += 1
+            used = 0
+        assert gpu_target < num_gpus
+        device_map[f'{layer_prefix}.layers.{i}'] = gpu_target
+        used += 1
+
+    return device_map
+
+def load_model(llm_device="cuda", device_map=None):
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path,trust_remote_code=True)
+    model_config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=True)
+    model = AutoModel.from_pretrained(model_name_or_path, config=model_config, trust_remote_code=True)
+    
+    if torch.cuda.is_available() and llm_device.lower().startswith("cuda"):
+        # 根据当前设备GPU数量决定是否进行多卡部署
+        num_gpus = torch.cuda.device_count()
+        if num_gpus < 2 and device_map is None:
+            model = model.half().cuda()
+        else:
+            from accelerate import dispatch_model
+
+            # model = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=True,
+            #                                   config=model_config, **kwargs)
+            # 可传入device_map自定义每张卡的部署情况
+            if device_map is None:
+                device_map = auto_configure_device_map(num_gpus, use_lora)
+
+            model = dispatch_model(
+                model.half(), device_map=device_map)
+    else:
+        model = model.float().to(llm_device)
+
+    model = model.eval()
+
+    return tokenizer, model

models/falcon_hf.py

@@ -0,0 +1,70 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
+
+max_token: int = 10000  # 10000 # 64
+temperature: float = 0.75
+top_p = 0.9
+use_lora = False
+
+# model_name_or_path = "Hannes-Epoch/falcon-7b-instruct-8bit" # not work, miss file
+
+
+def load_model(opt="gptq"):
+    if "pt" == opt:
+        return load_pt_model()
+    elif "gptq" == opt:
+        return load_gptq_model()
+    else:
+        raise Exception("not supported opt: {}".format(opt))
+
+########################################################################################################
+
+def load_gptq_model():
+    model_name_or_path = "TheBloke/falcon-7b-instruct-GPTQ"
+    # You could also download the model locally, and access it there
+    # model_name_or_path = "/path/to/TheBloke_falcon-7b-instruct-GPTQ"
+
+    model_basename = "gptq_model-4bit-64g"
+
+    use_triton = False
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name_or_path, use_fast=True)
+
+    model = AutoGPTQForCausalLM.from_quantized(model_name_or_path,
+                                               model_basename=model_basename,
+                                               use_safetensors=True,
+                                               trust_remote_code=True,
+                                               device="cuda:0",
+                                               use_triton=use_triton,
+                                               quantize_config=None)
+
+    return tokenizer, model
+
+
+########################################################################################################
+
+def load_pt_model():
+    model_name_or_path = "tiiuae/falcon-7b"
+    # model_name_or_path = "tiiuae/falcon-7b-instruct"
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name_or_path,
+        trust_remote_code=True,
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name_or_path,
+        trust_remote_code=True,
+        device_map='auto',
+        # load_in_8bit=True, # not working "RWForCausalLM.__init__() got an unexpected keyword argument 'load_in_8bit'"
+    )
+
+    return tokenizer, model
+
+########################################################################################################

models/vicuna_bin.py

@@ -0,0 +1,32 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+
+
+from llama_cpp import Llama, LlamaCache
+from common import torch_gc
+
+
+max_token: int = 10000
+temperature: float = 0.75
+top_p = 0.9
+
+def load_model():
+    model_name_or_path = "/root/下载/ggml-vic13b-q5_1.bin" 
+
+    params = {
+        'model_path': str(model_name_or_path),
+        'n_ctx': 2048,
+        'seed': 0,
+        'n_threads': 8,
+        'n_gpu_layers': 40,
+        'n_batch': 512,
+        'verbose': True,
+    }
+    model = Llama(**params)
+    model.set_cache(LlamaCache)
+
+    tokenizer = model.tokenizer()
+
+    return tokenizer, model

prompts/__init__.py

@@ -0,0 +1,4 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+

prompts/chatgpt_query.py

@@ -0,0 +1,134 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+def get_input_text(lang_opt, short_memory, input_paragraph, input_instruction, input_long_term_memory, new_character_prompt):
+    if "en" == lang_opt:
+        input_text = f"""I need you to help me write a novel. Now I give you a memory (a brief summary) of 400 words, you should use it to store the key content of what has been written so that you can keep track of very long context. For each time, I will give you your current memory (a brief summary of previous stories. You should use it to store the key content of what has been written so that you can keep track of very long context), the previously written paragraph, and instructions on what to write in the next paragraph. 
+        I need you to write:
+        1. Output Paragraph: the next paragraph of the novel. The output paragraph should contain around 20 sentences and should follow the input instructions.
+        2. Output Memory: The updated memory. You should first explain which sentences in the input memory are no longer necessary and why, and then explain what needs to be added into the memory and why. After that you should write the updated memory. The updated memory should be similar to the input memory except the parts you previously thought that should be deleted or added. The updated memory should only store key information. The updated memory should never exceed 20 sentences!
+        3. Output Instruction:  instructions of what to write next (after what you have written). You should output 3 different instructions, each is a possible interesting continuation of the story. Each output instruction should contain around 5 sentences
+        Here are the inputs: 
+
+        Input Memory:  
+        {short_memory}
+
+        Input Paragraph:
+        {input_paragraph}
+
+        Input Instruction:
+        {input_instruction}
+
+        Input Related Paragraphs:
+        {input_long_term_memory}
+        
+        Now start writing, organize your output by strictly following the output format as below:
+        Output Paragraph: 
+        <string of output paragraph>, around 20 sentences.
+
+        Output Memory: 
+        Rational: <string that explain how to update the memory>;
+        Updated Memory: <string of updated memory>, around 10 to 20 sentences
+
+        Output Instruction: 
+        Instruction 1: <content for instruction 1>, around 5 sentences
+        Instruction 2: <content for instruction 2>, around 5 sentences
+        Instruction 3: <content for instruction 3>, around 5 sentences
+
+        Very important!! The updated memory should only store key information. The updated memory should never contain over 500 words!
+        Finally, remember that you are writing a novel. Write like a novelist and do not move too fast when writing the output instructions for the next paragraph. Remember that the chapter will contain over 10 paragraphs and the novel will contain over 100 chapters. And this is just the beginning. Just write some interesting staffs that will happen next. Also, think about what plot can be attractive for common readers when writing output instructions. 
+
+        Very Important: 
+        You should first explain which sentences in the input memory are no longer necessary and why, and then explain what needs to be added into the memory and why. After that, you start rewrite the input memory to get the updated memory. 
+        {new_character_prompt}
+        """
+
+    elif "zh1" == lang_opt:
+        input_text = f"""I need you to help me write a novel. Now I give you a memory (a brief summary) of 400 words, you should use it to store the key content of what has been written so that you can keep track of very long context. For each time, I will give you your current memory (a brief summary of previous stories. You should use it to store the key content of what has been written so that you can keep track of very long context), the previously written paragraph, and instructions on what to write in the next paragraph. 
+        I need you to write:
+        1. Output Paragraph: the next paragraph of the novel. The output paragraph should contain around 20 sentences and should follow the input instructions.
+        2. Output Memory: The updated memory. You should first explain which sentences in the input memory are no longer necessary and why, and then explain what needs to be added into the memory and why. After that you should write the updated memory. The updated memory should be similar to the input memory except the parts you previously thought that should be deleted or added. The updated memory should only store key information. The updated memory should never exceed 20 sentences!
+        3. Output Instruction:  instructions of what to write next (after what you have written). You should output 3 different instructions, each is a possible interesting continuation of the story. Each output instruction should contain around 5 sentences
+        4. 非常重要！请将输出信息内容全部转化为中文，注意要符合中文母语的语法和���词习惯。
+        Here are the inputs: 
+
+        Input Memory:  
+        {short_memory}
+
+        Input Paragraph:
+        {input_paragraph}
+
+        Input Instruction:
+        {input_instruction}
+
+        Input Related Paragraphs:
+        {input_long_term_memory}
+        
+        Now start writing, organize your output by strictly following the output format as below:
+        Output Paragraph: 
+        <string of output paragraph>, around 20 sentences.
+
+        Output Memory: 
+        Rational: <string that explain how to update the memory>;
+        Updated Memory: <string of updated memory>, around 10 to 20 sentences
+
+        Output Instruction: 
+        Instruction 1: <content for instruction 1>, around 5 sentences
+        Instruction 2: <content for instruction 2>, around 5 sentences
+        Instruction 3: <content for instruction 3>, around 5 sentences
+
+        Very important!! The updated memory should only store key information. The updated memory should never contain over 500 words!
+        Finally, remember that you are writing a novel. Write like a novelist and do not move too fast when writing the output instructions for the next paragraph. Remember that the chapter will contain over 10 paragraphs and the novel will contain over 100 chapters. And this is just the beginning. Just write some interesting staffs that will happen next. Also, think about what plot can be attractive for common readers when writing output instructions. 
+
+        Very Important: 
+        You should first explain which sentences in the input memory are no longer necessary and why, and then explain what needs to be added into the memory and why. After that, you start rewrite the input memory to get the updated memory. 
+        非常重要！请将输出信息内容全部转化为中文，注意要符合中文母语的语法和用词习惯。
+        {new_character_prompt}
+        """
+
+    elif "zh2" == lang_opt:
+        input_text = f"""我需要你帮我写一部小说。现在我给你一个400字的记忆（一个简短的总结），你应该用它来存储已经写好的关键内容，这样你就可以记录很长的上下文。每一次，我都会给你当前的记忆（以前的故事的简要总结。你应该用它来存储所写内容的关键内容，这样你就能记下很长的上下文），之前写的段落，以及下一段要写的内容的指示。
+        我需要你来写：
+        1. 输出段落：小说的下一个段落。输出段应包含约20句话，并应遵循输入指示。
+        2. 输出记忆： 更新后的记忆。你应该首先解释输入记忆中的哪些句子不再需要，为什么，然后解释需要添加到记忆中的内容，为什么。之后，你应该写出更新的记忆。除了你之前认为应该删除或添加的部分，更新后的记忆应该与输入的记忆相似。更新后的记忆应该只存储关键信息。更新后的记忆不应该超过20个句子！
+        3. 输出指令：接下来要写什么的指令（在你写完之后）。你应该输出3个不同的指令，每个指令都是故事的一个可能的有趣的延续。每个输出指令应该包含大约5个句子
+        下面是输入的内容： 
+
+        输入内存：  
+        {short_memory}
+
+        输入段落：
+        {input_paragraph}
+
+        输入指令：
+        {input_instruction}。
+
+        输入相关段落：
+        {input_long_term_memory}
+        
+        现在开始写，严格按照下面的输出格式来组织你的输出：
+        输出段落： 
+        <输出段落的字符串>，大约20句话。
+
+        输出记忆： 
+        理性： <解释如何更新内存的字符串>；
+        更新的记忆： <更新内存的字符串>，大约10到20句话
+
+        输出指令： 
+        指令1：<指令1的内容>，大约5句话
+        指令2：<指令2的内容>，大约5句话
+        指令3：<指令3的内容>，大约5句话
+
+        非常重要！! 更新的内存应该只存储关键信息。更新后的记忆不应该包含超过500个字！！！！
+        最后，记住你在写一本小说。像小说家一样写作，在写下一段的输出指令时不要走得太快。记住，这一章将包含10多段，而小说将包含100多章。而这仅仅是个开始。就要写一些接下来会发生的有趣的职员。另外，在写输出说明时，要考虑什么情节能吸引普通读者。
+
+        非常重要： 
+        你应该首先解释输入存储器中的哪些句子不再需要，为什么，然后解释需要添加到存储器中的内容，为什么。之后，你开始重写输入内存，得到更新的内存。
+        {new_character_prompt}
+        """
+
+    else:
+        raise Exception("not supported lang_opt: {}".format(lang_opt))
+
+    return input_text

prompts/human_simulator.py

@@ -0,0 +1,120 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+
+def get_input_text(lang_opt, previous_paragraph, memory, writer_new_paragraph, user_edited_plan):
+    if "en" == lang_opt:
+        input_text = f"""
+        Now imagine you are a novelist writing a Chinese novel with the help of ChatGPT. You will be given a previously written paragraph (wrote by you), and a paragraph written by your ChatGPT assistant, a summary of the main storyline maintained by your ChatGPT assistant, and a plan of what to write next proposed by your ChatGPT assistant.
+        I need you to write:
+        1. Extended Paragraph: Extend the new paragraph written by the ChatGPT assistant to twice the length of the paragraph written by your ChatGPT assistant.
+        2. Selected Plan: Copy the plan proposed by your ChatGPT assistant.
+        3. Revised Plan: Revise the selected plan into an outline of the next paragraph.
+        
+        Previously written paragraph:  
+        {previous_paragraph}
+
+        The summary of the main storyline maintained by your ChatGPT assistant:
+        {memory}
+
+        The new paragraph written by your ChatGPT assistant:
+        {writer_new_paragraph}
+
+        The plan of what to write next proposed by your ChatGPT assistant:
+        {user_edited_plan}
+
+        Now start writing, organize your output by strictly following the output format as below,所有输出仍然保持是中文:
+        
+        Extended Paragraph: 
+        <string of output paragraph>, around 40-50 sentences.
+
+        Selected Plan: 
+        <copy the plan here>
+
+        Revised Plan:
+        <string of revised plan>, keep it short, around 5-7 sentences.
+
+        Very Important:
+        Remember that you are writing a novel. Write like a novelist and do not move too fast when writing the plan for the next paragraph. Think about how the plan can be attractive for common readers when selecting and extending the plan. Remember to follow the length constraints! Remember that the chapter will contain over 10 paragraphs and the novel will contain over 100 chapters. And the next paragraph will be the second paragraph of the second chapter. You need to leave space for future stories.
+
+        """
+
+    elif "zh1" == lang_opt:
+        input_text = f"""
+        Now imagine you are a novelist writing a Chinese novel with the help of ChatGPT. You will be given a previously written paragraph (wrote by you), and a paragraph written by your ChatGPT assistant, a summary of the main storyline maintained by your ChatGPT assistant, and a plan of what to write next proposed by your ChatGPT assistant.
+        I need you to write:
+        1. Extended Paragraph: Extend the new paragraph written by the ChatGPT assistant to twice the length of the paragraph written by your ChatGPT assistant.
+        2. Selected Plan: Copy the plan proposed by your ChatGPT assistant.
+        3. Revised Plan: Revise the selected plan into an outline of the next paragraph.
+        4. 非常重要！请将输出信息内容全部转化为中文，注意要符合中文母语的语法和用词习惯。
+        
+        Previously written paragraph:  
+        {previous_paragraph}
+
+        The summary of the main storyline maintained by your ChatGPT assistant:
+        {memory}
+
+        The new paragraph written by your ChatGPT assistant:
+        {writer_new_paragraph}
+
+        The plan of what to write next proposed by your ChatGPT assistant:
+        {user_edited_plan}
+
+        Now start writing, organize your output by strictly following the output format as below,所有输出仍然保持是中文:
+        
+        Extended Paragraph: 
+        <string of output paragraph>, around 40-50 sentences.
+
+        Selected Plan: 
+        <copy the plan here>
+
+        Revised Plan:
+        <string of revised plan>, keep it short, around 5-7 sentences.
+
+        Very Important:
+        Remember that you are writing a novel. Write like a novelist and do not move too fast when writing the plan for the next paragraph. Think about how the plan can be attractive for common readers when selecting and extending the plan. Remember to follow the length constraints! Remember that the chapter will contain over 10 paragraphs and the novel will contain over 100 chapters. And the next paragraph will be the second paragraph of the second chapter. You need to leave space for future stories.
+        非常重要！请将输出信息内容全部转化为中文，注意要符合中文母语的语法和用词习惯。
+        
+        """
+
+    elif "zh2" == lang_opt:
+        input_text = f"""
+        现在想象一下，你是一个小说家，在ChatGPT的帮助下写一本中文小说。你会得到一个先前写好的段落（由你写），和一个由你的ChatGPT助手写的段落，一个由你的ChatGPT助手保持的主要故事情节的总结，以及一个由你的ChatGPT助手提出的下一步写作计划。
+        我需要你写：
+        1. 扩展段落： 将ChatGPT助手写的新段落延长到你的ChatGPT助手���写段落的两倍。
+        2. 选定计划： 复制您的ChatGPT助手提出的计划。
+        3. 修订的计划： 将选定的计划修改为下一段的纲要。
+        
+        以前写的段落： 
+        {previous_paragraph}
+
+        由你的ChatGPT助手维护的主要故事情节的摘要：
+        {memory}
+
+        您的ChatGPT助理写的新段落：
+        {writer_new_paragraph}
+
+        您的ChatGPT助理提出的下一步写作计划：
+        {user_edited_plan}
+
+        现在开始写，严格按照下面的输出格式来组织你的输出，所有输出仍然保持是中文：
+
+        扩展段落： 
+        <输出段落的字符串>，大约40-50个句子。
+
+        选定的计划： 
+        <在此复制计划>
+
+        修改后的计划：
+        <修改后的计划字符串>，保持简短，大约5-7句话。
+
+        非常重要：
+        记住你在写一本小说。像小说家一样写作，在写下一段的计划时不要走得太快。在选择和扩展计划时，要考虑计划如何对普通读者具有吸引力。记住要遵循长度限制! 记住，这一章将包含10多段，而小说将包含100多章。而下一段将是第二章的第二段。你需要为未来的故事留出空间。
+        
+        """
+
+    else:
+        raise Exception("not supported lang_opt: {}".format(lang_opt))
+
+    return input_text

prompts/llm_query.py

@@ -0,0 +1,135 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+
+def get_input_text(lang_opt, short_memory, input_paragraph, input_instruction, input_long_term_memory, new_character_prompt):
+    if "en" == lang_opt:
+        input_text = f"""I need you to help me write a novel. Now I give you a memory (a brief summary) of 400 words, you should use it to store the key content of what has been written so that you can keep track of very long context. For each time, I will give you your current memory (a brief summary of previous stories. You should use it to store the key content of what has been written so that you can keep track of very long context), the previously written paragraph, and instructions on what to write in the next paragraph. 
+        I need you to write:
+        1. Output Paragraph: the next paragraph of the novel. The output paragraph should contain around 20 sentences and should follow the input instructions.
+        2. Output Memory: The updated memory. You should first explain which sentences in the input memory are no longer necessary and why, and then explain what needs to be added into the memory and why. After that you should write the updated memory. The updated memory should be similar to the input memory except the parts you previously thought that should be deleted or added. The updated memory should only store key information. The updated memory should never exceed 20 sentences!
+        3. Output Instruction:  instructions of what to write next (after what you have written). You should output 3 different instructions, each is a possible interesting continuation of the story. Each output instruction should contain around 5 sentences
+        Here are the inputs: 
+
+        Input Memory:  
+        {short_memory}
+
+        Input Paragraph:
+        {input_paragraph}
+
+        Input Instruction:
+        {input_instruction}
+
+        Input Related Paragraphs:
+        {input_long_term_memory}
+        
+        Now start writing, organize your output by strictly following the output format as below:
+        Output Paragraph: 
+        <string of output paragraph>, around 20 sentences.
+
+        Output Memory: 
+        Rational: <string that explain how to update the memory>;
+        Updated Memory: <string of updated memory>, around 10 to 20 sentences
+
+        Output Instruction: 
+        Instruction 1: <content for instruction 1>, around 5 sentences
+        Instruction 2: <content for instruction 2>, around 5 sentences
+        Instruction 3: <content for instruction 3>, around 5 sentences
+
+        Very important!! The updated memory should only store key information. The updated memory should never contain over 500 words!
+        Finally, remember that you are writing a novel. Write like a novelist and do not move too fast when writing the output instructions for the next paragraph. Remember that the chapter will contain over 10 paragraphs and the novel will contain over 100 chapters. And this is just the beginning. Just write some interesting staffs that will happen next. Also, think about what plot can be attractive for common readers when writing output instructions. 
+
+        Very Important: 
+        You should first explain which sentences in the input memory are no longer necessary and why, and then explain what needs to be added into the memory and why. After that, you start rewrite the input memory to get the updated memory. 
+        {new_character_prompt}
+        """
+
+    elif "zh1" == lang_opt:
+        input_text = f"""I need you to help me write a novel. Now I give you a memory (a brief summary) of 400 words, you should use it to store the key content of what has been written so that you can keep track of very long context. For each time, I will give you your current memory (a brief summary of previous stories. You should use it to store the key content of what has been written so that you can keep track of very long context), the previously written paragraph, and instructions on what to write in the next paragraph. 
+        I need you to write:
+        1. Output Paragraph: the next paragraph of the novel. The output paragraph should contain around 20 sentences and should follow the input instructions.
+        2. Output Memory: The updated memory. You should first explain which sentences in the input memory are no longer necessary and why, and then explain what needs to be added into the memory and why. After that you should write the updated memory. The updated memory should be similar to the input memory except the parts you previously thought that should be deleted or added. The updated memory should only store key information. The updated memory should never exceed 20 sentences!
+        3. Output Instruction:  instructions of what to write next (after what you have written). You should output 3 different instructions, each is a possible interesting continuation of the story. Each output instruction should contain around 5 sentences
+        4. 非常重要！请将输出信息内容全部转化为中文，注意要符合中文母语的语法��用词习惯。
+        Here are the inputs: 
+
+        Input Memory:  
+        {short_memory}
+
+        Input Paragraph:
+        {input_paragraph}
+
+        Input Instruction:
+        {input_instruction}
+
+        Input Related Paragraphs:
+        {input_long_term_memory}
+        
+        Now start writing, organize your output by strictly following the output format as below:
+        Output Paragraph: 
+        <string of output paragraph>, around 20 sentences.
+
+        Output Memory: 
+        Rational: <string that explain how to update the memory>;
+        Updated Memory: <string of updated memory>, around 10 to 20 sentences
+
+        Output Instruction: 
+        Instruction 1: <content for instruction 1>, around 5 sentences
+        Instruction 2: <content for instruction 2>, around 5 sentences
+        Instruction 3: <content for instruction 3>, around 5 sentences
+
+        Very important!! The updated memory should only store key information. The updated memory should never contain over 500 words!
+        Finally, remember that you are writing a novel. Write like a novelist and do not move too fast when writing the output instructions for the next paragraph. Remember that the chapter will contain over 10 paragraphs and the novel will contain over 100 chapters. And this is just the beginning. Just write some interesting staffs that will happen next. Also, think about what plot can be attractive for common readers when writing output instructions. 
+
+        Very Important: 
+        You should first explain which sentences in the input memory are no longer necessary and why, and then explain what needs to be added into the memory and why. After that, you start rewrite the input memory to get the updated memory. 
+        非常重要！请将输出信息内容全部转化为中文，注意要符合中文母语的语法和用词习惯。
+        {new_character_prompt}
+        """
+
+    elif "zh2" == lang_opt:
+        input_text = f"""我需要你帮我写一部小说。现在我给你一个400字的记忆（一个简短的总结），你应该用它来存储已经写好的关键内容，这样你就可以记录很长的上下文。每一次，我都会给你当前的记忆（以前的故事的简要总结。你应该用它来存储所写内容的关键内容，这样你就能记下很长的上下文），之前写的段落，以及下一段要写的内容的指示。
+        我需要你来写：
+        1. 输出段落：小说的下一个段落。输出段应包含约20句话，并应遵循输入指示。
+        2. 输出记忆： 更新后的记忆。你应该首先解释输入记忆中的哪些句子不再需要，为什么，然后解释需要添加到记忆中的内容，为什么。之后，你应该写出更新的记忆。除了你之前认为应该删除或添加的部分，更新后的记忆应该与输入的记忆相似。更新后的记忆应该只存储关键信息。更新后的记忆不应该超过20个句子！
+        3. 输出指令：接下来要写什么的指令（在你写完之后）。你应该输出3个不同的指令，每个指令都是故事的一个可能的有趣的延续。每个输出指令应该包含大约5个句子
+        下面是输入的内容： 
+
+        输入内存：  
+        {short_memory}
+
+        输入段落：
+        {input_paragraph}
+
+        输入指令：
+        {input_instruction}。
+
+        输入相关段落：
+        {input_long_term_memory}
+        
+        现在开始写，严格按照下面的输出格式来组织你的输出：
+        输出段落： 
+        <输出段落的字符串>，大约20句话。
+
+        输出记忆： 
+        理性： <解释如何更新内存的字符串>；
+        更新的记忆： <更新内存的字符串>，大约10到20句话
+
+        输出指令： 
+        指令1：<指令1的内容>，大约5句话
+        指令2：<指令2的内容>，大约5句话
+        指令3：<指令3的内容>，大约5句话
+
+        非常重要！! 更新的内存应该只存储关键信息。更新后的记忆不应该包含超过500个字！！！！
+        最后，记住你在写一本小说。像小说家一样写作，在写下一段的输出指令时不要走得太快。记住，这一章将包含10多段，而小说将包含100多章。而这仅仅是个开始。就要写一些接下来会发生的有趣的职员。另外，在写输出说明时，要考虑什么情节能吸引普通读者。
+
+        非常重要： 
+        你应该首先解释输入存储器中的哪些句子不再需要，为什么，然后解释需要添加到存储器中的内容，为什么。之后，你开始重写输入内存，得到更新的内存。
+        {new_character_prompt}
+        """
+
+    else:
+        raise Exception("not supported lang_opt: {}".format(lang_opt))
+            
+    return input_text

prompts/service_init.py

@@ -0,0 +1,79 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+
+def get_init_prompt(lang_opt, novel_type, description):
+    if "en" == lang_opt:
+        return f"""
+        Please write a {novel_type} novel{description} with 50 chapters. Follow the format below precisely:
+
+        Begin with the name of the novel.
+        Next, write an outline for the first chapter. The outline should describe the background and the beginning of the novel.
+        Write the first three paragraphs with their indication of the novel based on your outline. Write in a novelistic style and take your time to set the scene.
+        Write a summary that captures the key information of the three paragraphs.
+        Finally, write three different instructions for what to write next, each containing around five sentences. Each instruction should present a possible, interesting continuation of the story.
+        The output format should follow these guidelines:
+        Name: <name of the novel>
+        Outline: <outline for the first chapter>
+        Paragraph 1: <content for paragraph 1>
+        Paragraph 2: <content for paragraph 2>
+        Paragraph 3: <content for paragraph 3>
+        Summary: <content of summary>
+        Instruction 1: <content for instruction 1>
+        Instruction 2: <content for instruction 2>
+        Instruction 3: <content for instruction 3>
+
+        Make sure to be precise and follow the output format strictly.
+
+        """
+    elif "zh1" == lang_opt:
+        return f"""
+        Please write a {novel_type} novel{description} with 50 chapters. Follow the format below precisely:
+
+        Begin with the name of the novel.
+        Next, write an outline for the first chapter. The outline should describe the background and the beginning of the novel.
+        Write the first three paragraphs with their indication of the novel based on your outline. Write in a novelistic style and take your time to set the scene.
+        Write a summary that captures the key information of the three paragraphs.
+        Finally, write three different instructions for what to write next, each containing around five sentences. Each instruction should present a possible, interesting continuation of the story.
+        The output format should follow these guidelines:
+        名称： <name of the novel>
+        概述： <outline for the first chapter>
+        段落1： <content for paragraph 1>
+        段落2： <content for paragraph 2>
+        段落3： <content for paragraph 3>
+        总结： <content of summary>
+        指令1： <content for instruction 1>
+        指令2： <content for instruction 2>
+        指令3：<content for instruction 3>
+
+        Make sure to be precise and follow the output format strictly.
+        非常重要！请将输出信息内容全部转化为中文，注意要符合中文母语的语法和用词习惯。
+
+        """
+    
+    elif "zh2" == lang_opt:
+        return f"""
+        请写一篇{novel_type}的小说{description}，有50个章节。准确遵循以下格式：
+
+        以小说的名称开始。
+        接下来，写出第一章的大纲。大纲应描述小说的背景和开头。
+        根据你的提纲写出前三段，并说明小说的内容。用小说的风格来写，慢慢地设置场景。
+        写一个摘要，抓住这三段的关键信息。
+        最后，写出三个不同的指示，说明接下来要写什么，每个指示包含大约五句话。每个指示都应该提出一个可能的、有趣的故事的延续。
+        输出格式应遵循这些准则：
+        名称： <小说的名称>
+        概述： <第一章的大纲>
+        段落1： <第1段的内容>
+        段落2： <第2段的内容>
+        段落3： <第3段的内容>
+        总结： <摘要的内容>。
+        指令1： <指令1的内容>
+        指令2： <指令2的内容>
+        指令3：<指令3的内容>
+
+        请务必准确无误，并严格遵守输出格式。
+        """
+
+    else:
+        raise Exception(f"not supported language: {lang_opt}")

recurrent_llm.py

@@ -0,0 +1,121 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+import torch
+import random
+from sentence_transformers import util
+
+from utils import get_content_between_a_b
+from prompts.llm_query import get_input_text
+from global_config import lang_opt, llm_model_opt
+
+if "openai" == llm_model_opt:
+    from utils.openai_util import get_api_response
+elif "vicuna" == llm_model_opt:
+    from utils.vicuna_util import get_api_response
+elif "chatglm" == llm_model_opt:
+    from utils.chatglm_util import get_api_response
+elif "baichuan" == llm_model_opt:
+    from utils.baichuan_util import get_api_response
+elif "aquila" == llm_model_opt:
+    from utils.aquila_util import get_api_response
+elif "falcon" == llm_model_opt:
+    from utils.falcon_util import get_api_response
+else:
+    raise Exception("not supported llm model name: {}".format(llm_model_opt))
+
+
+class RecurrentLLM:
+
+    def __init__(self, input, short_memory, long_memory, memory_index, embedder, model, tokenizer):
+        print("AIWriter loaded by RecurrentLLM")
+        self.input = input
+        self.short_memory = short_memory
+        self.long_memory = long_memory
+        self.embedder = embedder
+        self.model = model
+        self.tokenizer = tokenizer
+        if self.long_memory and not memory_index:
+            self.memory_index = self.embedder.encode(
+                self.long_memory, convert_to_tensor=True)
+        self.output = {}
+
+    def prepare_input(self, new_character_prob=0.1, top_k=2):
+
+        input_paragraph = self.input["output_paragraph"]
+        input_instruction = self.input["output_instruction"]
+
+        instruction_embedding = self.embedder.encode(
+            input_instruction, convert_to_tensor=True)
+
+        # get the top 3 most similar paragraphs from memory
+
+        memory_scores = util.cos_sim(
+            instruction_embedding, self.memory_index)[0]
+        top_k_idx = torch.topk(memory_scores, k=top_k)[1]
+        top_k_memory = [self.long_memory[idx] for idx in top_k_idx]
+        # combine the top 3 paragraphs
+        input_long_term_memory = '\n'.join(
+            [f"Related Paragraphs {i+1} :" + selected_memory for i, selected_memory in enumerate(top_k_memory)])
+        # randomly decide if a new character should be introduced
+        if random.random() < new_character_prob:
+            new_character_prompt = f"If it is reasonable, you can introduce a new character in the output paragrah and add it into the memory."
+        else:
+            new_character_prompt = ""
+
+        input_text = get_input_text(lang_opt, self.short_memory, input_paragraph, input_instruction, input_long_term_memory, new_character_prompt)
+
+        return input_text
+
+    def parse_output(self, output):
+        try:
+            output_paragraph = get_content_between_a_b(
+                'Output Paragraph:', 'Output Memory', output)
+            output_memory_updated = get_content_between_a_b(
+                'Updated Memory:', 'Output Instruction:', output)
+            self.short_memory = output_memory_updated
+            ins_1 = get_content_between_a_b(
+                'Instruction 1:', 'Instruction 2', output)
+            ins_2 = get_content_between_a_b(
+                'Instruction 2:', 'Instruction 3', output)
+            lines = output.splitlines()
+            # content of Instruction 3 may be in the same line with I3 or in the next line
+            if lines[-1] != '\n' and lines[-1].startswith('Instruction 3'):
+                ins_3 = lines[-1][len("Instruction 3:"):]
+            elif lines[-1] != '\n':
+                ins_3 = lines[-1]
+
+            output_instructions = [ins_1, ins_2, ins_3]
+            assert len(output_instructions) == 3
+
+            output = {
+                "input_paragraph": self.input["output_paragraph"],
+                "output_memory": output_memory_updated,  # feed to human
+                "output_paragraph": output_paragraph,
+                "output_instruction": [instruction.strip() for instruction in output_instructions]
+            }
+
+            return output
+        except:
+            return None
+
+    def step(self, response_file=None):
+
+        prompt = self.prepare_input()
+
+        print(prompt+'\n'+'\n')
+
+        response = get_api_response(self.model, self.tokenizer, prompt)
+
+        self.output = self.parse_output(response)
+        while self.output == None:
+            response = get_api_response(self.model, self.tokenizer, prompt)
+            self.output = self.parse_output(response)
+        if response_file:
+            with open(response_file, 'a', encoding='utf-8') as f:
+                f.write(f"Writer's output here:\n{response}\n\n")
+
+        self.long_memory.append(self.input["output_paragraph"])
+        self.memory_index = self.embedder.encode(
+            self.long_memory, convert_to_tensor=True)

recurrentgpt.py

@@ -0,0 +1,104 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+import torch
+import random
+from sentence_transformers import util
+
+from utils import get_content_between_a_b, get_api_response
+from prompts.chatgpt_query import get_input_text
+from global_config import lang_opt
+
+
+class RecurrentGPT:
+
+    def __init__(self, input, short_memory, long_memory, memory_index, embedder):
+        print("AIWriter loaded by RecurrentGPT")
+        self.input = input
+        self.short_memory = short_memory
+        self.long_memory = long_memory
+        self.embedder = embedder
+        if self.long_memory and not memory_index:
+            self.memory_index = self.embedder.encode(
+                self.long_memory, convert_to_tensor=True)
+        self.output = {}
+
+    def prepare_input(self, new_character_prob=0.1, top_k=2):
+
+        input_paragraph = self.input["output_paragraph"]
+        input_instruction = self.input["output_instruction"]
+
+        instruction_embedding = self.embedder.encode(
+            input_instruction, convert_to_tensor=True)
+
+        # get the top 3 most similar paragraphs from memory
+
+        memory_scores = util.cos_sim(
+            instruction_embedding, self.memory_index)[0]
+        top_k_idx = torch.topk(memory_scores, k=top_k)[1]
+        top_k_memory = [self.long_memory[idx] for idx in top_k_idx]
+        # combine the top 3 paragraphs
+        input_long_term_memory = '\n'.join(
+            [f"Related Paragraphs {i+1} :" + selected_memory for i, selected_memory in enumerate(top_k_memory)])
+        # randomly decide if a new character should be introduced
+        if random.random() < new_character_prob:
+            new_character_prompt = f"If it is reasonable, you can introduce a new character in the output paragrah and add it into the memory."
+        else:
+            new_character_prompt = ""
+
+        input_text = get_input_text(lang_opt, self.short_memory, input_paragraph, input_instruction, input_long_term_memory, new_character_prompt)
+        
+        return input_text
+
+    def parse_output(self, output):
+        try:
+            output_paragraph = get_content_between_a_b(
+                'Output Paragraph:', 'Output Memory', output)
+            output_memory_updated = get_content_between_a_b(
+                'Updated Memory:', 'Output Instruction:', output)
+            self.short_memory = output_memory_updated
+            ins_1 = get_content_between_a_b(
+                'Instruction 1:', 'Instruction 2', output)
+            ins_2 = get_content_between_a_b(
+                'Instruction 2:', 'Instruction 3', output)
+            lines = output.splitlines()
+            # content of Instruction 3 may be in the same line with I3 or in the next line
+            if lines[-1] != '\n' and lines[-1].startswith('Instruction 3'):
+                ins_3 = lines[-1][len("Instruction 3:"):]
+            elif lines[-1] != '\n':
+                ins_3 = lines[-1]
+
+            output_instructions = [ins_1, ins_2, ins_3]
+            assert len(output_instructions) == 3
+
+            output = {
+                "input_paragraph": self.input["output_paragraph"],
+                "output_memory": output_memory_updated,  # feed to human
+                "output_paragraph": output_paragraph,
+                "output_instruction": [instruction.strip() for instruction in output_instructions]
+            }
+
+            return output
+        except:
+            return None
+
+    def step(self, response_file=None):
+
+        prompt = self.prepare_input()
+
+        print(prompt+'\n'+'\n')
+
+        response = get_api_response(prompt)
+
+        self.output = self.parse_output(response)
+        while self.output == None:
+            response = get_api_response(prompt)
+            self.output = self.parse_output(response)
+        if response_file:
+            with open(response_file, 'a', encoding='utf-8') as f:
+                f.write(f"Writer's output here:\n{response}\n\n")
+
+        self.long_memory.append(self.input["output_paragraph"])
+        self.memory_index = self.embedder.encode(
+            self.long_memory, convert_to_tensor=True)

requirements.txt

@@ -0,0 +1,19 @@
+pip install -U sentence-transformers
+pip install --upgrade openai
+
+
+pip install bitsandbytes==0.39.0
+pip install transformers@git+https://github.com/huggingface/transformers.git
+pip install peft@git+https://github.com/huggingface/peft.git
+pip install accelerate@git+https://github.com/huggingface/accelerate.git
+
+
+llama-cpp-python@git+https://github.com/abetlen/llama-cpp-python.git
+
+
+pip install -U flagai
+pip install bminf
+
+
+pip install auto-gptq
+pip install einops

utils/__init__.py

@@ -0,0 +1,181 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+import re
+from global_config import lang_opt, llm_model_opt
+
+if "openai" == llm_model_opt:
+    from utils.openai_util import get_api_response
+elif "vicuna" == llm_model_opt:
+    from utils.vicuna_util import get_api_response
+elif "chatglm" == llm_model_opt:
+    from utils.chatglm_util import get_api_response
+elif "baichuan" == llm_model_opt:
+    from utils.baichuan_util import get_api_response
+elif "aquila" == llm_model_opt:
+    from utils.aquila_util import get_api_response
+elif "falcon" == llm_model_opt:
+    from utils.falcon_util import get_api_response
+else:
+    raise Exception("not supported llm model name: {}".format(llm_model_opt))
+
+
+def get_content_between_a_b(a, b, text):
+    if "en" == lang_opt:
+        if "vicuna" == llm_model_opt:
+            return re.search(f"{a}(.*?)\n(.*?){b}", text, re.DOTALL).group(1).strip()
+        elif "openai" == llm_model_opt:
+            return re.search(f"{a}(.*?)\n{b}", text, re.DOTALL).group(1).strip()
+        elif llm_model_opt in ["chatglm", "baichuan", "aquila", "falcon"]:
+            return re.search(f"{a}(.*?)\n(.*?){b}", text, re.DOTALL).group(1).strip()
+        else:
+            raise Exception(
+                "not supported llm model name: {}".format(llm_model_opt))
+
+    elif lang_opt in ["zh1", "zh2"]:
+        if "vicuna" == llm_model_opt:
+            match = re.search(f"{a}(.*?)\n(.*?){b}", text, re.DOTALL)
+        elif "openai" == llm_model_opt:
+            match = re.search(f"{a}(.*?)\n{b}", text, re.DOTALL)
+        elif llm_model_opt in ["chatglm", "baichuan", "aquila", "falcon"]:
+            match = re.search(f"{a}(.*?)\n(.*?){b}", text, re.DOTALL)
+        else:
+            raise Exception(
+                "not supported llm model name: {}".format(llm_model_opt))
+
+        if match:
+            return match.group(1).strip()
+        else:
+            if "1" in a or "2" in a or "3" in a:
+                a = ''.join(a.split(" "))
+            if "1" in b or "2" in b or "3" in b:
+                b = "".join(b.split(" "))
+
+            if "vicuna" == llm_model_opt:
+                match = re.search(f"{a}(.*?)\n(.*?){b}", text, re.DOTALL)
+            elif "openai" == llm_model_opt:
+                match = re.search(f"{a}(.*?)\n{b}", text, re.DOTALL)
+            elif llm_model_opt in ["chatglm", "baichuan", "aquila", "falcon"]:
+                match = re.search(f"{a}(.*?)\n(.*?){b}", text, re.DOTALL)
+            else:
+                raise Exception(
+                    "not supported llm model name: {}".format(llm_model_opt))
+
+            if match:
+                return match.group(1).strip()
+            else:
+                # 处理找不到匹配内容的情况
+                return "翻译时出现错误请重试"  # 或者返回其他默认值或采取其他的处理方式
+    else:
+        raise Exception(f"not supported language: {lang_opt}")
+
+
+def get_init(init_text=None, text=None, response_file=None, model=None, tokenizer=None):
+    """
+    init_text: if the title, outline, and the first 3 paragraphs are given in a .txt file, directly read
+    text: if no .txt file is given, use init prompt to generate
+    """
+    if not init_text:
+        response = get_api_response(model, tokenizer, text)
+        print("response: {}".format(response))
+
+        if response_file:
+            with open(response_file, 'a', encoding='utf-8') as f:
+                f.write(f"Init output here:\n{response}\n\n")
+    else:
+        with open(init_text, 'r', encoding='utf-8') as f:
+            response = f.read()
+        f.close()
+    paragraphs = {
+        "name": "",
+        "Outline": "",
+        "Paragraph 1": "",
+        "Paragraph 2": "",
+        "Paragraph 3": "",
+        "Summary": "",
+        "Instruction 1": "",
+        "Instruction 2": "",
+        "Instruction 3": ""
+    }
+
+    if "en" == lang_opt:
+        paragraphs['name'] = get_content_between_a_b(
+            'Name:', 'Outline', response)
+
+        paragraphs['Paragraph 1'] = get_content_between_a_b(
+            'Paragraph 1:', 'Paragraph 2:', response)
+        paragraphs['Paragraph 2'] = get_content_between_a_b(
+            'Paragraph 2:', 'Paragraph 3:', response)
+        paragraphs['Paragraph 3'] = get_content_between_a_b(
+            'Paragraph 3:', 'Summary', response)
+        paragraphs['Summary'] = get_content_between_a_b(
+            'Summary:', 'Instruction 1', response)
+        paragraphs['Instruction 1'] = get_content_between_a_b(
+            'Instruction 1:', 'Instruction 2', response)
+        paragraphs['Instruction 2'] = get_content_between_a_b(
+            'Instruction 2:', 'Instruction 3', response)
+        lines = response.splitlines()
+        # content of Instruction 3 may be in the same line with I3 or in the next line
+        if lines[-1] != '\n' and lines[-1].startswith('Instruction 3'):
+            paragraphs['Instruction 3'] = lines[-1][len("Instruction 3:"):]
+        elif lines[-1] != '\n':
+            paragraphs['Instruction 3'] = lines[-1]
+        # Sometimes it gives Chapter outline, sometimes it doesn't
+        for line in lines:
+            if line.startswith('Chapter'):
+                paragraphs['Outline'] = get_content_between_a_b(
+                    'Outline:', 'Chapter', response)
+                break
+        if paragraphs['Outline'] == '':
+            paragraphs['Outline'] = get_content_between_a_b(
+                'Outline:', 'Paragraph', response)
+
+    elif lang_opt in ["zh1", "zh2"]:
+        paragraphs['name'] = get_content_between_a_b('名称：', '概述：', response)
+
+        paragraphs['Paragraph 1'] = get_content_between_a_b(
+            '段落 1：', '段落 2：', response)
+        paragraphs['Paragraph 2'] = get_content_between_a_b(
+            '段落 2：', '段落 3：', response)
+        paragraphs['Paragraph 3'] = get_content_between_a_b(
+            '段落 3：', '总结：', response)
+        paragraphs['Summary'] = get_content_between_a_b(
+            '总结：', '指令 1', response)
+        paragraphs['Instruction 1'] = get_content_between_a_b(
+            '指令 1：', '指令 2：', response)
+        paragraphs['Instruction 2'] = get_content_between_a_b(
+            '指令 2：', '指令 3：', response)
+        lines = response.splitlines()
+        # content of Instruction 3 may be in the same line with I3 or in the next line
+        if lines[-1] != '\n' and lines[-1].startswith('Instruction 3'):
+            paragraphs['Instruction 3'] = lines[-1][len("Instruction 3:"):]
+        elif lines[-1] != '\n':
+            paragraphs['Instruction 3'] = lines[-1]
+        # Sometimes it gives Chapter outline, sometimes it doesn't
+        for line in lines:
+            if line.startswith('Chapter'):
+                paragraphs['Outline'] = get_content_between_a_b(
+                    '概述：', 'Chapter', response)
+                break
+        if paragraphs['Outline'] == '':
+            paragraphs['Outline'] = get_content_between_a_b(
+                '概述：', '段落', response)
+
+    return paragraphs
+
+
+def get_chatgpt_response(model, prompt):
+    response = ""
+    for data in model.ask(prompt):
+        response = data["message"]
+    model.delete_conversation(model.conversation_id)
+    model.reset_chat()
+    return response
+
+
+def parse_instructions(instructions):
+    output = ""
+    for i in range(len(instructions)):
+        output += f"{i+1}. {instructions[i]}\n"
+    return output

utils/aquila_util.py

@@ -0,0 +1,84 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+
+import torch
+from flagai.model.predictor.predictor import Predictor
+from flagai.model.predictor.aquila import aquila_generate
+from models.aquila_fa import max_token, temperature, top_p
+from common import torch_gc
+from global_config import lang_opt
+
+# for Aquila on FlagAI
+def get_api_response(model, tokenizer, content: str, max_tokens=None):
+
+    if "en" == lang_opt:
+        system_role_content = 'You are a helpful and creative assistant for writing novel.'
+    elif "zh1" == lang_opt:
+        system_role_content = 'You are a helpful and creative assistant for writing novel.\
+                You are must always in Chinese.重要，你需要使用中文与我进行交流。'
+    elif "zh2" == lang_opt:
+        system_role_content = '你是写小说的好帮手，有创意的助手。'
+    else:
+        raise Exception(f"not supported language: {lang_opt}")
+
+    print("===> Question:")
+    print(content)
+    print("<==="+"="*100)
+
+    predictor = Predictor(model, tokenizer)
+    content = f'{content}'
+    with torch.no_grad():
+        out = predictor.predict_generate_randomsample(
+            content, out_max_length=max_token, temperature=temperature, top_p=top_p)
+        response = out
+
+    torch_gc()
+
+    print("===> Generated Text: ")
+    print(response)
+    print("<==="+"="*100)
+
+    return response
+
+# # for Aquila on HuggingFace
+# def get_api_response(model, tokenizer, content: str, max_tokens=None):
+
+#     if "en" == lang_opt:
+#         system_role_content = 'You are a helpful and creative assistant for writing novel.'
+#     elif "zh1" == lang_opt:
+#         system_role_content = 'You are a helpful and creative assistant for writing novel.\
+#                 You are must always in Chinese.重要，你需要使用中文与我进行交流。'
+#     elif "zh2" == lang_opt:
+#         system_role_content = '你是写小说的好帮手，有创意的助手。'
+#     else:
+#         raise Exception(f"not supported language: {lang_opt}")
+
+#     print("===> Question:")
+#     print(content)
+#     print("<==="+"="*100)
+
+#     with torch.no_grad():
+#         ret = model.generate(
+#             **tokenizer(content, return_tensors='pt').to('cuda'),
+#             do_sample=False,
+#             max_new_tokens=max_token,
+#             temperature=temperature,
+#             top_p=top_p,
+#             use_cache=True
+#         )
+#         output_ids = ret[0].detach().cpu().numpy().tolist()
+#         if 100007 in output_ids:
+#             output_ids = output_ids[:output_ids.index(100007)]
+#         elif 0 in output_ids:
+#             output_ids = output_ids[:output_ids.index(0)]
+#         response = tokenizer.decode(output_ids)
+
+#     torch_gc()
+
+#     print("===> Generated Text: ")
+#     print(response)
+#     print("<==="+"="*100)
+
+#     return response

utils/baichuan_util.py

@@ -0,0 +1,53 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+from transformers import TextStreamer
+
+from models.baichuan_hf import max_token, temperature, top_p
+from common import torch_gc
+from global_config import lang_opt
+
+
+def get_api_response(model, tokenizer, content: str, max_tokens=None):
+
+    if "en" == lang_opt:
+        system_role_content = 'You are a helpful and creative assistant for writing novel.'
+    elif "zh1" == lang_opt:
+        system_role_content = 'You are a helpful and creative assistant for writing novel.\
+                You are must always in Chinese.重要，你需要使用中文与我进行交流。'
+    elif "zh2" == lang_opt:
+        system_role_content = '你是写小说的好帮手，有创意的助手。'
+    else:
+        raise Exception(f"not supported language: {lang_opt}")
+
+    print("===> Question:")
+    print(content)
+    print("<==="+"="*100)
+
+    streamer = TextStreamer(tokenizer, 
+                            skip_prompt=True,
+                            skip_special_tokens=True
+                            )
+
+    # inputs = tokenizer(content, return_tensors='pt')
+    inputs = tokenizer("<human>:{}\n<bot>:".format(content), return_tensors='pt')
+    # inputs = inputs.to('cuda') # UserWarning: You are calling .generate() with the `input_ids` being on a device type different than your model's device. `input_ids` is on cuda, whereas the model is on cpu. You may experience unexpected behaviors or slower generation. Please make sure that you have put `input_ids` to the correct device by calling for example input_ids = input_ids.to('cpu') before running `.generate()`.
+    inputs = inputs.to('cpu')
+    generate_ids = model.generate(**inputs,
+                                  max_new_tokens=max_token,
+                                  top_p=top_p,
+                                  temperature=temperature,
+                                  repetition_penalty=1.1,
+                                  streamer=streamer,
+                                  )
+    response = tokenizer.decode(
+        generate_ids.cpu()[0], skip_special_tokens=True)
+
+    torch_gc()
+
+    print("===> Generated Text: ")
+    print(response)
+    print("<==="+"="*100)
+
+    return response

utils/chatglm_util.py

@@ -0,0 +1,42 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+
+from models.chatglm_hf import max_token, temperature, top_p
+from common import torch_gc
+from global_config import lang_opt
+
+
+def get_api_response(model, tokenizer, content: str, max_tokens=None):
+
+    if "en" == lang_opt:
+        system_role_content = 'You are a helpful and creative assistant for writing novel.'
+    elif "zh1" == lang_opt:
+        system_role_content = 'You are a helpful and creative assistant for writing novel.\
+                You are must always in Chinese.重要，你需要使用中文与我进行交流。'
+    elif "zh2" == lang_opt:
+        system_role_content = '你是写小说的好帮手，有创意的助手。'
+    else:
+        raise Exception(f"not supported language: {lang_opt}")
+
+    print("===> Question:")
+    print(content)
+    print("<==="+"="*100)
+
+    response, history = model.chat(
+        tokenizer,
+        content,
+        history=[],
+        max_length=max_token,
+        temperature=temperature,
+        top_p=top_p,
+    )
+
+    torch_gc()
+
+    print("===> Generated Text: ")
+    print(response)
+    print("<==="+"="*100)
+
+    return response

utils/falcon_util.py

@@ -0,0 +1,46 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+from models.falcon_hf import max_token, temperature, top_p
+from common import torch_gc
+from global_config import lang_opt
+
+
+def get_api_response(model, tokenizer, content: str, max_tokens=None):
+
+    if "en" == lang_opt:
+        system_role_content = 'You are a helpful and creative assistant for writing novel.'
+    elif "zh1" == lang_opt:
+        system_role_content = 'You are a helpful and creative assistant for writing novel.\
+                You are must always in Chinese.重要，你需要使用中文与我进行交流。'
+    elif "zh2" == lang_opt:
+        system_role_content = '你是写小说的好帮手，有创意的助手。'
+    else:
+        raise Exception(f"not supported language: {lang_opt}")
+
+    print("===> Question:")
+    print(content)
+    print("<==="+"="*100)
+
+    inputs = tokenizer(content,
+                       return_tensors='pt',
+                       return_token_type_ids=False, #  ValueError: The following model_kwargs are not used by the model: ['token_type_ids'] (note: typos in the generate arguments will also show up in this list)
+                       )
+    inputs = inputs.to('cuda:0')
+    output = model.generate(**inputs,
+                          max_new_tokens=max_token,
+                          top_p=top_p,
+                          temperature=temperature,
+                          repetition_penalty=1.1,
+                        #   eos_token_id=tokenizer.eos_token_id,
+                          )
+    response = tokenizer.decode(output.cpu()[0], skip_special_tokens=True)
+
+    torch_gc()
+
+    print("===> Generated Text: ")
+    print(response)
+    print("<==="+"="*100)
+
+    return response

utils/openai_util.py

@@ -0,0 +1,38 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+
+import openai
+
+from global_config import lang_opt
+
+
+def get_api_response(model, tokenizer, content: str, max_tokens=None):
+
+    if "en" == lang_opt:
+        system_role_content = 'You are a helpful and creative assistant for writing novel.'
+    elif "zh1" == lang_opt:
+        system_role_content = 'You are a helpful and creative assistant for writing novel.\
+                You are must always in Chinese.重要，你需要使用中文与我进行交流。'
+    elif "zh2" == lang_opt:
+        system_role_content = '你是写小说的好帮手，有创意的助手。'
+    else:
+        raise Exception(f"not supported language: {lang_opt}")
+
+    response = openai.ChatCompletion.create(
+        model='gpt-3.5-turbo',
+        messages=[{
+            'role': 'system',
+            'content': system_role_content
+        }, {
+            'role': 'user',
+            'content': content,
+        }],
+        temperature=0.5,
+        max_tokens=max_tokens
+    )
+
+    return response['choices'][0]['message']['content']
+
+

utils/vicuna_util.py

@@ -0,0 +1,59 @@
+#!python
+# -*- coding: utf-8 -*-
+# @author: Kun
+
+
+from models.vicuna_bin import max_token, temperature, top_p
+from common import torch_gc
+from global_config import lang_opt
+
+
+def get_api_response(model, tokenizer, content: str, max_tokens=None):
+
+    if "en" == lang_opt:
+        system_role_content = 'You are a helpful and creative assistant for writing novel.'
+    elif "zh1" == lang_opt:
+        system_role_content = 'You are a helpful and creative assistant for writing novel.\
+                You are must always in Chinese.重要，你需要使用中文与我进行交流。'
+    elif "zh2" == lang_opt:
+        system_role_content = '你是写小说的好帮手，有创意的助手。'
+    else:
+        raise Exception(f"not supported language: {lang_opt}")
+
+    print("===> Question:")
+    print(content)
+    print("<==="+"="*100)
+
+    content = content.encode()
+    tokens = model.tokenize(content)
+
+    output = b""
+    count = 0
+    token_count = 10000
+    top_k = 40
+    repetition_penalty = 1.1
+    for token in model.generate(tokens,
+                                top_k=top_k,
+                                top_p=top_p,
+                                temp=temperature,
+                                repeat_penalty=repetition_penalty):
+        text = model.detokenize([token])
+        # print(text)
+        output += text
+
+        count += 1
+        if count >= token_count or (token == model.token_eos()):
+            break
+
+    response = output.decode()
+    # print("===> [vicuna][generate] response: {}".format(response))
+
+    torch_gc()
+
+    print("===> Generated Text: ")
+    print(response)
+    print("<==="+"="*100)
+
+    return response
+
+