From 1c80e9a09ed9fce0809e2b893d946306df0f9f7c Mon Sep 17 00:00:00 2001
From: hiyouga <hiyouga@buaa.edu.cn>
Date: Thu, 16 Nov 2023 17:12:23 +0800
Subject: [PATCH] fix web ui demo

---
 src/llmtuner/__init__.py      |  2 +-
 src/llmtuner/webui/chatter.py | 10 ++++++++--
 src/llmtuner/webui/engine.py  |  3 ++-
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/src/llmtuner/__init__.py b/src/llmtuner/__init__.py
index e4376f9c..712be7c2 100644
--- a/src/llmtuner/__init__.py
+++ b/src/llmtuner/__init__.py
@@ -7,4 +7,4 @@ from llmtuner.train import export_model, run_exp
 from llmtuner.webui import create_ui, create_web_demo
 
 
-__version__ = "0.3.0"
+__version__ = "0.3.1"
diff --git a/src/llmtuner/webui/chatter.py b/src/llmtuner/webui/chatter.py
index 92f4bcb1..5b27aaa8 100644
--- a/src/llmtuner/webui/chatter.py
+++ b/src/llmtuner/webui/chatter.py
@@ -29,8 +29,14 @@ class WebChatModel(ChatModel):
         if not lazy_init: # read arguments from command line
             super().__init__()
 
-        if demo_mode: # load openchat 3.5 by default
-            super().__init__(dict(model_name_or_path="openchat/openchat_3.5", template="openchat"))
+        if demo_mode: # load config.json by default
+            import json
+            try:
+                with open("demo_config.json", "r", encoding="utf-8") as f:
+                    args = json.load(f)
+            except:
+                raise ValueError("Cannot find `demo_config.json` at current directory.")
+            super().__init__(args)
 
     @property
     def loaded(self) -> bool:
diff --git a/src/llmtuner/webui/engine.py b/src/llmtuner/webui/engine.py
index 3e9f077d..bdac09dd 100644
--- a/src/llmtuner/webui/engine.py
+++ b/src/llmtuner/webui/engine.py
@@ -13,6 +13,7 @@ from llmtuner.webui.utils import get_time
 class Engine:
 
     def __init__(self, demo_mode: Optional[bool] = False, pure_chat: Optional[bool] = False) -> None:
+        self.demo_mode = demo_mode
         self.pure_chat = pure_chat
         self.manager = Manager()
         self.runner = Runner(self.manager, demo_mode=demo_mode)
@@ -22,7 +23,7 @@ class Engine:
         return {self.manager.get_elem_by_name(k): gr.update(**v) for k, v in resume_dict.items()}
 
     def resume(self) -> Generator[Dict[Component, Dict[str, Any]], None, None]:
-        user_config = load_config()
+        user_config = load_config() if not self.demo_mode else {}
         lang = user_config.get("lang", None) or "en"
 
         init_dict = {