release v1.7.7

CURRENT_CHANGE.md

@@ -1,25 +1,7 @@
 ## Changes
 
-### Features
+- avoid program lag caused by frequent triggering of read/write operations due to Linux file system notification
-
+- improve styles
-- add Docker support (#291) @LonghronShen
-
-### Fixes
-
-- fix a generation exception caused by potentially dangerous regex being passed into the stop array
-- fix max_tokens parameter of Chat page not being passed to backend
-- fix the issue where penalty_decay and global_penalty are not being passed to the backend default config when running
-  the model through client
-
-### Improvements
-
-- prevent 'torch' has no attribute 'cuda' error in torch_gc, so user can use CPU or WebGPU (#302)
-
-### Chores
-
-- bump dependencies
-- add pre-release workflow
-- dep_check.py now ignores GPUtil
 
 ## Install
 

README.md

@@ -94,7 +94,8 @@ English | [简体中文](README_ZH.md) | [日本語](README_JA.md)
 - Built-in model conversion tool.
 - Built-in download management and remote model inspection.
 - Built-in one-click LoRA Finetune. (Windows Only)
-- Can also be used as an OpenAI ChatGPT and GPT-Playground client. (Fill in the API URL and API Key in Settings page)
+- Can also be used as an OpenAI ChatGPT, GPT-Playground, Ollama and more clients. (Fill in the API URL and API Key in
+  Settings page)
 - Multilingual localization.
 - Theme switching.
 - Automatic updates.

README_JA.md

@@ -89,8 +89,8 @@
 - 内蔵モデル変換ツール
 - ダウンロード管理とリモートモデル検査機能内蔵
 - 内蔵のLoRA微調整機能を搭載しています (Windowsのみ)
-- このプログラムは、OpenAI ChatGPTとGPT Playgroundのクライアントとしても使用できます（設定ページで `API URL` と `API Key`
+- このプログラムは、OpenAI ChatGPT、GPT Playground、Ollama などのクライアントとしても使用できます（設定ページで `API URL`
-  を入力してください）
+  と `API Key` を入力してください）
 - 多言語ローカライズ
 - テーマ切り替え
 - 自動アップデート

README_ZH.md

@@ -83,7 +83,7 @@ API兼容的接口，这意味着一切ChatGPT客户端都是RWKV客户端。
 - 内置模型转换工具
 - 内置下载管理和远程模型检视
 - 内置一键LoRA微调 (仅限Windows)
-- 也可用作 OpenAI ChatGPT 和 GPT Playground 客户端 (在设置内填写API URL和API Key)
+- 也可用作 OpenAI ChatGPT, GPT Playground, Ollama 等服务的客户端 (在设置内填写API URL和API Key)
 - 多语言本地化
 - 主题切换
 - 自动更新

backend-golang/app.go

@@ -7,7 +7,11 @@ import (
 	"context"
 	"errors"
 	"io"
+	"log"
+	"net"
 	"net/http"
+	"net/http/httputil"
+	"net/url"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -27,6 +31,7 @@ type App struct {
 	HasConfigData bool
 	ConfigData    map[string]any
 	Dev           bool
+	proxyPort     int
 	exDir         string
 	cmdPrefix     string
 }
@@ -36,6 +41,63 @@ func NewApp() *App {
 	return &App{}
 }
 
+func (a *App) newFetchProxy() {
+	go func() {
+		handler := func(w http.ResponseWriter, r *http.Request) {
+			if r.Method == "OPTIONS" {
+				w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
+				w.Header().Set("Access-Control-Allow-Headers", "*")
+				w.Header().Set("Access-Control-Allow-Origin", "*")
+				return
+			}
+			proxy := &httputil.ReverseProxy{
+				ModifyResponse: func(res *http.Response) error {
+					res.Header.Set("Access-Control-Allow-Origin", "*")
+					return nil
+				},
+				Director: func(req *http.Request) {
+					realTarget := req.Header.Get("Real-Target")
+					if realTarget != "" {
+						realTarget, err := url.PathUnescape(realTarget)
+						if err != nil {
+							log.Printf("Error decoding target URL: %v\n", err)
+							return
+						}
+						target, err := url.Parse(realTarget)
+						if err != nil {
+							log.Printf("Error parsing target URL: %v\n", err)
+							return
+						}
+						req.Header.Set("Accept", "*/*")
+						req.Header.Del("Origin")
+						req.Header.Del("Referer")
+						req.Header.Del("Real-Target")
+						req.Header.Del("Sec-Fetch-Dest")
+						req.Header.Del("Sec-Fetch-Mode")
+						req.Header.Del("Sec-Fetch-Site")
+						req.URL.Scheme = target.Scheme
+						req.URL.Host = target.Host
+						req.URL.Path = target.Path
+						req.URL.RawQuery = url.PathEscape(target.RawQuery)
+						log.Println("Proxying to", realTarget)
+					} else {
+						log.Println("Real-Target header is missing")
+					}
+				},
+			}
+			proxy.ServeHTTP(w, r)
+		}
+		http.HandleFunc("/", handler)
+		listener, err := net.Listen("tcp", "127.0.0.1:0")
+		if err != nil {
+			return
+		}
+		a.proxyPort = listener.Addr().(*net.TCPAddr).Port
+
+		http.Serve(listener, nil)
+	}()
+}
+
 // startup is called when the app starts. The context is saved
 // so we can call the runtime methods
 func (a *App) OnStartup(ctx context.Context) {
@@ -76,6 +138,7 @@ func (a *App) OnStartup(ctx context.Context) {
 	a.midiLoop()
 	a.watchFs()
 	a.monitorHardware()
+	a.newFetchProxy()
 }
 
 func (a *App) OnBeforeClose(ctx context.Context) bool {
@@ -239,3 +302,7 @@ func (a *App) RestartApp() error {
 func (a *App) GetPlatform() string {
 	return runtime.GOOS
 }
+
+func (a *App) GetProxyPort() int {
+	return a.proxyPort
+}

backend-golang/rwkv.go

@@ -227,12 +227,12 @@ func (a *App) InstallPyDep(python string, cnMirror bool) (string, error) {
 
 	if runtime.GOOS == "windows" {
 		ChangeFileLine("./py310/python310._pth", 3, "Lib\\site-packages")
-		installScript := python + " ./backend-python/get-pip.py -i https://pypi.tuna.tsinghua.edu.cn/simple --no-warn-script-location\n" +
+		installScript := python + " ./backend-python/get-pip.py -i https://mirrors.aliyun.com/pypi/simple --no-warn-script-location\n" +
 			python + " -m pip install torch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 --index-url https://download.pytorch.org/whl/cu117 --no-warn-script-location\n" +
-			python + " -m pip install -r ./backend-python/requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple --no-warn-script-location\n" +
+			python + " -m pip install -r ./backend-python/requirements.txt -i https://mirrors.aliyun.com/pypi/simple --no-warn-script-location\n" +
 			"exit"
 		if !cnMirror {
-			installScript = strings.Replace(installScript, " -i https://pypi.tuna.tsinghua.edu.cn/simple", "", -1)
+			installScript = strings.Replace(installScript, " -i https://mirrors.aliyun.com/pypi/simple", "", -1)
 		}
 		err = os.WriteFile(a.exDir+"install-py-dep.bat", []byte(installScript), 0644)
 		if err != nil {
@@ -242,7 +242,7 @@ func (a *App) InstallPyDep(python string, cnMirror bool) (string, error) {
 	}
 
 	if cnMirror {
-		return Cmd(python, "-m", "pip", "install", "-r", "./backend-python/requirements_without_cyac.txt", "-i", "https://pypi.tuna.tsinghua.edu.cn/simple")
+		return Cmd(python, "-m", "pip", "install", "-r", "./backend-python/requirements_without_cyac.txt", "-i", "https://mirrors.aliyun.com/pypi/simple")
 	} else {
 		return Cmd(python, "-m", "pip", "install", "-r", "./backend-python/requirements_without_cyac.txt")
 	}

backend-python/routes/completion.py

@@ -53,6 +53,9 @@ class ChatCompletionBody(ModelConfigBody):
     assistant_name: Union[str, None] = Field(
         None, description="Internal assistant name", min_length=1
     )
+    system_name: Union[str, None] = Field(
+        None, description="Internal system name", min_length=1
+    )
     presystem: bool = Field(
         True, description="Whether to insert default system prompt at the beginning"
     )
@@ -68,6 +71,7 @@ class ChatCompletionBody(ModelConfigBody):
                 "stop": None,
                 "user_name": None,
                 "assistant_name": None,
+                "system_name": None,
                 "presystem": True,
                 "max_tokens": 1000,
                 "temperature": 1,
@@ -252,20 +256,9 @@ async def eval_rwkv(
                 }
 
 
-@router.post("/v1/chat/completions", tags=["Completions"])
+def chat_template_old(
-@router.post("/chat/completions", tags=["Completions"])
+    model: TextRWKV, body: ChatCompletionBody, interface: str, user: str, bot: str
-async def chat_completions(body: ChatCompletionBody, request: Request):
+):
-    model: TextRWKV = global_var.get(global_var.Model)
-    if model is None:
-        raise HTTPException(status.HTTP_400_BAD_REQUEST, "model not loaded")
-
-    if body.messages is None or body.messages == []:
-        raise HTTPException(status.HTTP_400_BAD_REQUEST, "messages not found")
-
-    interface = model.interface
-    user = model.user if body.user_name is None else body.user_name
-    bot = model.bot if body.assistant_name is None else body.assistant_name
-
     is_raven = model.rwkv_type == RWKVType.Raven
 
     completion_text: str = ""
@@ -334,6 +327,53 @@ The following is a coherent verbose detailed conversation between a girl named {
         completion_text += append_message + "\n\n"
     completion_text += f"{bot}{interface}"
 
+    return completion_text
+
+
+def chat_template(
+    model: TextRWKV, body: ChatCompletionBody, interface: str, user: str, bot: str
+):
+    completion_text: str = ""
+    if body.presystem:
+        completion_text = (
+            f"{user}{interface} hi\n\n{bot}{interface} Hi. "
+            + "I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.\n\n"
+        )
+
+    system = "System" if body.system_name is None else body.system_name
+    for message in body.messages:
+        append_message: str = ""
+        if message.role == Role.User:
+            append_message = f"{user}{interface} " + message.content
+        elif message.role == Role.Assistant:
+            append_message = f"{bot}{interface} " + message.content
+        elif message.role == Role.System:
+            append_message = f"{system}{interface} " + message.content
+        completion_text += append_message + "\n\n"
+    completion_text += f"{bot}{interface}"
+
+    return completion_text
+
+
+@router.post("/v1/chat/completions", tags=["Completions"])
+@router.post("/chat/completions", tags=["Completions"])
+async def chat_completions(body: ChatCompletionBody, request: Request):
+    model: TextRWKV = global_var.get(global_var.Model)
+    if model is None:
+        raise HTTPException(status.HTTP_400_BAD_REQUEST, "model not loaded")
+
+    if body.messages is None or body.messages == []:
+        raise HTTPException(status.HTTP_400_BAD_REQUEST, "messages not found")
+
+    interface = model.interface
+    user = model.user if body.user_name is None else body.user_name
+    bot = model.bot if body.assistant_name is None else body.assistant_name
+
+    if model.version < 5:
+        completion_text = chat_template_old(model, body, interface, user, bot)
+    else:
+        completion_text = chat_template(model, body, interface, user, bot)
+
     user_code = model.pipeline.decode([model.pipeline.encode(user)[0]])
     bot_code = model.pipeline.decode([model.pipeline.encode(bot)[0]])
     if type(body.stop) == str:
@@ -343,8 +383,8 @@ The following is a coherent verbose detailed conversation between a girl named {
         body.stop.append(f"\n\n{bot_code}")
     elif body.stop is None:
         body.stop = default_stop
-    if not body.presystem:
+    # if not body.presystem:
-        body.stop.append("\n\n")
+    #     body.stop.append("\n\n")
 
     if body.stream:
         return EventSourceResponse(

backend-python/rwkv_pip/cpp/model.py

@@ -9,6 +9,9 @@ class RWKV:
         self.model = rwkv_cpp_model.RWKVModel(self.library, model_path)
         self.w = {}  # fake weight
         self.w["emb.weight"] = [0] * self.model.n_vocab
+        self.version = (
+            self.model.arch_version_major + self.model.arch_version_minor / 10
+        )
 
     def forward(self, tokens: List[int], state: Union[Any, None] = None):
         return self.model.eval_sequence_in_chunks(tokens, state, use_numpy=True)

backend-python/rwkv_pip/cpp/rwkv_cpp_model.py

@@ -52,9 +52,14 @@ class RWKVModel:
         if 'gpu_layers_count' in kwargs:
             gpu_layer_count = kwargs['gpu_layers_count']
 
-        assert os.path.isfile(model_path), f'{model_path} is not a file'
+        if not os.path.isfile(model_path):
-        assert thread_count > 0, 'Thread count must be > 0'
+            raise ValueError(f'{model_path} is not a file')
-        assert gpu_layer_count >= 0, 'GPU layer count must be >= 0'
+        
+        if not (thread_count > 0):
+            raise ValueError('Thread count must be > 0')  
+
+        if not (gpu_layer_count >= 0):
+            raise ValueError('GPU layer count must be >= 0')
 
         self._library: rwkv_cpp_shared_library.RWKVSharedLibrary = shared_library
 
@@ -84,10 +89,19 @@ class RWKVModel:
             Count of layers to offload onto the GPU, must be >= 0.
         """
 
-        assert layer_count >= 0, 'Layer count must be >= 0'
+        if not (layer_count >= 0):
+            raise ValueError('Layer count must be >= 0')
 
         return self._library.rwkv_gpu_offload_layers(self._ctx, layer_count)
 
+    @property
+    def arch_version_major(self) -> int:
+        return self._library.rwkv_get_arch_version_major(self._ctx)
+
+    @property
+    def arch_version_minor(self) -> int:
+        return self._library.rwkv_get_arch_version_minor(self._ctx)
+
     @property
     def n_vocab(self) -> int:
         return self._library.rwkv_get_n_vocab(self._ctx)
@@ -133,7 +147,8 @@ class RWKVModel:
             Logits vector of shape (n_vocab); state for the next step.
         """
 
-        assert self._valid, 'Model was freed'
+        if not self._valid:
+            raise ValueError('Model was freed')
 
         use_numpy = self._detect_numpy_usage([state_in, state_out, logits_out], use_numpy)
 
@@ -207,7 +222,8 @@ class RWKVModel:
             Logits vector of shape (n_vocab); state for the next step.
         """
 
-        assert self._valid, 'Model was freed'
+        if not self._valid:
+            raise ValueError('Model was freed')
 
         use_numpy = self._detect_numpy_usage([state_in, state_out, logits_out], use_numpy)
 
@@ -281,7 +297,8 @@ class RWKVModel:
             Logits vector of shape (n_vocab); state for the next step.
         """
 
-        assert self._valid, 'Model was freed'
+        if not self._valid:
+            raise ValueError('Model was freed')
 
         use_numpy = self._detect_numpy_usage([state_in, state_out, logits_out], use_numpy)
 
@@ -320,7 +337,8 @@ class RWKVModel:
         The object must not be used anymore after calling this method.
         """
 
-        assert self._valid, 'Already freed'
+        if not self._valid:
+            raise ValueError('Already freed')
 
         self._valid = False
 
@@ -344,16 +362,25 @@ class RWKVModel:
     def _validate_tensor(self, tensor: NumpyArrayOrPyTorchTensor, name: str, size: int) -> None:
         if self._is_pytorch_tensor(tensor):
             tensor: torch.Tensor = tensor
-            assert tensor.device == torch.device('cpu'), f'{name} is not on CPU'
+            
-            assert tensor.dtype == torch.float32, f'{name} is not of type float32'
+            if tensor.device != torch.device('cpu'):
-            assert tensor.shape == (size,), f'{name} has invalid shape {tensor.shape}, expected ({size})'
+                raise ValueError(f'{name} is not on CPU')
-            assert tensor.is_contiguous(), f'{name} is not contiguous'
+            if tensor.dtype != torch.float32:
+                raise ValueError(f'{name} is not of type float32')
+            if tensor.shape != (size,):
+                raise ValueError(f'{name} has invalid shape {tensor.shape}, expected ({size})')
+            if not tensor.is_contiguous():
+                raise ValueError(f'{name} is not contiguous')
         else:
             import numpy as np
             tensor: np.ndarray = tensor
-            assert tensor.dtype == np.float32, f'{name} is not of type float32'
+
-            assert tensor.shape == (size,), f'{name} has invalid shape {tensor.shape}, expected ({size})'
+            if tensor.dtype != np.float32:
-            assert tensor.data.contiguous, f'{name} is not contiguous'
+                raise ValueError(f'{name} is not of type float32')
+            if tensor.shape != (size,):
+                raise ValueError(f'{name} has invalid shape {tensor.shape}, expected ({size})')
+            if not tensor.data.contiguous:
+                raise ValueError(f'{name} is not contiguous')
 
     def _get_data_ptr(self, tensor: NumpyArrayOrPyTorchTensor):
         if self._is_pytorch_tensor(tensor):

backend-python/rwkv_pip/cpp/rwkv_cpp_shared_library.py

@@ -6,21 +6,22 @@ import platform
 from typing import Optional, List, Tuple, Callable
 
 QUANTIZED_FORMAT_NAMES: Tuple[str, str, str, str, str] = (
-    'Q4_0',
+    "Q4_0",
-    'Q4_1',
+    "Q4_1",
-    'Q5_0',
+    "Q5_0",
-    'Q5_1',
+    "Q5_1",
-    'Q8_0'
+    "Q8_0",
 )
 
 P_FLOAT = ctypes.POINTER(ctypes.c_float)
 P_INT = ctypes.POINTER(ctypes.c_int32)
 
-class RWKVContext:
 
+class RWKVContext:
     def __init__(self, ptr: ctypes.pointer) -> None:
         self.ptr: ctypes.pointer = ptr
 
+
 class RWKVSharedLibrary:
     """
     Python wrapper around rwkv.cpp shared library.
@@ -39,7 +40,7 @@ class RWKVSharedLibrary:
         #  When Python is greater than 3.8, we need to reprocess the custom dll
         #  according to the documentation to prevent loading failure errors.
         #  https://docs.python.org/3/whatsnew/3.8.html#ctypes
-        if platform.system().lower() == 'windows':
+        if platform.system().lower() == "windows":
             self.library = ctypes.CDLL(shared_library_path, winmode=0)
         else:
             self.library = ctypes.cdll.LoadLibrary(shared_library_path)
@@ -47,39 +48,48 @@ class RWKVSharedLibrary:
         self.library.rwkv_init_from_file.argtypes = [ctypes.c_char_p, ctypes.c_uint32]
         self.library.rwkv_init_from_file.restype = ctypes.c_void_p
 
-        self.library.rwkv_gpu_offload_layers.argtypes = [ctypes.c_void_p, ctypes.c_uint32]
+        self.library.rwkv_gpu_offload_layers.argtypes = [
+            ctypes.c_void_p,
+            ctypes.c_uint32,
+        ]
         self.library.rwkv_gpu_offload_layers.restype = ctypes.c_bool
 
         self.library.rwkv_eval.argtypes = [
-            ctypes.c_void_p, # ctx
+            ctypes.c_void_p,  # ctx
-            ctypes.c_int32, # token
+            ctypes.c_int32,  # token
-            P_FLOAT, # state_in
+            P_FLOAT,  # state_in
-            P_FLOAT, # state_out
+            P_FLOAT,  # state_out
-            P_FLOAT  # logits_out
+            P_FLOAT,  # logits_out
         ]
         self.library.rwkv_eval.restype = ctypes.c_bool
 
         self.library.rwkv_eval_sequence.argtypes = [
-            ctypes.c_void_p, # ctx
+            ctypes.c_void_p,  # ctx
-            P_INT, # tokens
+            P_INT,  # tokens
-            ctypes.c_size_t, # token count
+            ctypes.c_size_t,  # token count
-            P_FLOAT, # state_in
+            P_FLOAT,  # state_in
-            P_FLOAT, # state_out
+            P_FLOAT,  # state_out
-            P_FLOAT  # logits_out
+            P_FLOAT,  # logits_out
         ]
         self.library.rwkv_eval_sequence.restype = ctypes.c_bool
 
         self.library.rwkv_eval_sequence_in_chunks.argtypes = [
-            ctypes.c_void_p, # ctx
+            ctypes.c_void_p,  # ctx
-            P_INT, # tokens
+            P_INT,  # tokens
-            ctypes.c_size_t, # token count
+            ctypes.c_size_t,  # token count
-            ctypes.c_size_t, # chunk size
+            ctypes.c_size_t,  # chunk size
-            P_FLOAT, # state_in
+            P_FLOAT,  # state_in
-            P_FLOAT, # state_out
+            P_FLOAT,  # state_out
-            P_FLOAT  # logits_out
+            P_FLOAT,  # logits_out
         ]
         self.library.rwkv_eval_sequence_in_chunks.restype = ctypes.c_bool
 
+        self.library.rwkv_get_arch_version_major.argtypes = [ctypes.c_void_p]
+        self.library.rwkv_get_arch_version_major.restype = ctypes.c_uint32
+
+        self.library.rwkv_get_arch_version_minor.argtypes = [ctypes.c_void_p]
+        self.library.rwkv_get_arch_version_minor.restype = ctypes.c_uint32
+
         self.library.rwkv_get_n_vocab.argtypes = [ctypes.c_void_p]
         self.library.rwkv_get_n_vocab.restype = ctypes.c_size_t
 
@@ -101,7 +111,11 @@ class RWKVSharedLibrary:
         self.library.rwkv_free.argtypes = [ctypes.c_void_p]
         self.library.rwkv_free.restype = None
 
-        self.library.rwkv_quantize_model_file.argtypes = [ctypes.c_char_p, ctypes.c_char_p, ctypes.c_char_p]
+        self.library.rwkv_quantize_model_file.argtypes = [
+            ctypes.c_char_p,
+            ctypes.c_char_p,
+            ctypes.c_char_p,
+        ]
         self.library.rwkv_quantize_model_file.restype = ctypes.c_bool
 
         self.library.rwkv_get_system_info_string.argtypes = []
@@ -109,7 +123,9 @@ class RWKVSharedLibrary:
 
         self.nullptr = ctypes.cast(0, ctypes.c_void_p)
 
-    def rwkv_init_from_file(self, model_file_path: str, thread_count: int) -> RWKVContext:
+    def rwkv_init_from_file(
+        self, model_file_path: str, thread_count: int
+    ) -> RWKVContext:
         """
         Loads the model from a file and prepares it for inference.
         Throws an exception in case of any error. Error messages would be printed to stderr.
@@ -122,9 +138,12 @@ class RWKVSharedLibrary:
             Count of threads to use, must be positive.
         """
 
-        ptr = self.library.rwkv_init_from_file(model_file_path.encode('utf-8'), ctypes.c_uint32(thread_count))
+        ptr = self.library.rwkv_init_from_file(
+            model_file_path.encode("utf-8"), ctypes.c_uint32(thread_count)
+        )
 
-        assert ptr is not None, 'rwkv_init_from_file failed, check stderr'
+        if ptr is None:
+            raise ValueError("rwkv_init_from_file failed, check stderr")
 
         return RWKVContext(ptr)
 
@@ -145,17 +164,20 @@ class RWKVSharedLibrary:
             Count of layers to offload onto the GPU, must be >= 0.
         """
 
-        assert layer_count >= 0, 'Layer count must be >= 0'
+        if not (layer_count >= 0):
+            raise ValueError("Layer count must be >= 0")
 
-        return self.library.rwkv_gpu_offload_layers(ctx.ptr, ctypes.c_uint32(layer_count))
+        return self.library.rwkv_gpu_offload_layers(
+            ctx.ptr, ctypes.c_uint32(layer_count)
+        )
 
     def rwkv_eval(
-            self,
+        self,
-            ctx: RWKVContext,
+        ctx: RWKVContext,
-            token: int,
+        token: int,
-            state_in_address: Optional[int],
+        state_in_address: Optional[int],
-            state_out_address: int,
+        state_out_address: int,
-            logits_out_address: int
+        logits_out_address: int,
     ) -> None:
         """
         Evaluates the model for a single token.
@@ -176,21 +198,22 @@ class RWKVSharedLibrary:
             Address of the first element of a FP32 buffer of size rwkv_get_logits_buffer_element_count. This buffer will be written to.
         """
 
-        assert self.library.rwkv_eval(
+        if not self.library.rwkv_eval(
             ctx.ptr,
             ctypes.c_int32(token),
             ctypes.cast(0 if state_in_address is None else state_in_address, P_FLOAT),
             ctypes.cast(state_out_address, P_FLOAT),
-            ctypes.cast(logits_out_address, P_FLOAT)
+            ctypes.cast(logits_out_address, P_FLOAT),
-        ), 'rwkv_eval failed, check stderr'
+        ):
+            raise ValueError("rwkv_eval failed, check stderr")
 
     def rwkv_eval_sequence(
-            self,
+        self,
-            ctx: RWKVContext,
+        ctx: RWKVContext,
-            tokens: List[int],
+        tokens: List[int],
-            state_in_address: Optional[int],
+        state_in_address: Optional[int],
-            state_out_address: int,
+        state_out_address: int,
-            logits_out_address: int
+        logits_out_address: int,
     ) -> None:
         """
         Evaluates the model for a sequence of tokens.
@@ -223,23 +246,24 @@ class RWKVSharedLibrary:
             Address of the first element of a FP32 buffer of size rwkv_get_logits_buffer_element_count. This buffer will be written to.
         """
 
-        assert self.library.rwkv_eval_sequence(
+        if not self.library.rwkv_eval_sequence(
             ctx.ptr,
             ctypes.cast((ctypes.c_int32 * len(tokens))(*tokens), P_INT),
             ctypes.c_size_t(len(tokens)),
             ctypes.cast(0 if state_in_address is None else state_in_address, P_FLOAT),
             ctypes.cast(state_out_address, P_FLOAT),
-            ctypes.cast(logits_out_address, P_FLOAT)
+            ctypes.cast(logits_out_address, P_FLOAT),
-        ), 'rwkv_eval_sequence failed, check stderr'
+        ):
+            raise ValueError("rwkv_eval_sequence failed, check stderr")
 
     def rwkv_eval_sequence_in_chunks(
-            self,
+        self,
-            ctx: RWKVContext,
+        ctx: RWKVContext,
-            tokens: List[int],
+        tokens: List[int],
-            chunk_size: int,
+        chunk_size: int,
-            state_in_address: Optional[int],
+        state_in_address: Optional[int],
-            state_out_address: int,
+        state_out_address: int,
-            logits_out_address: int
+        logits_out_address: int,
     ) -> None:
         """
         Evaluates the model for a sequence of tokens using `rwkv_eval_sequence`, splitting a potentially long sequence into fixed-length chunks.
@@ -269,15 +293,40 @@ class RWKVSharedLibrary:
             Address of the first element of a FP32 buffer of size rwkv_get_logits_buffer_element_count. This buffer will be written to.
         """
 
-        assert self.library.rwkv_eval_sequence_in_chunks(
+        if not self.library.rwkv_eval_sequence_in_chunks(
             ctx.ptr,
             ctypes.cast((ctypes.c_int32 * len(tokens))(*tokens), P_INT),
             ctypes.c_size_t(len(tokens)),
             ctypes.c_size_t(chunk_size),
             ctypes.cast(0 if state_in_address is None else state_in_address, P_FLOAT),
             ctypes.cast(state_out_address, P_FLOAT),
-            ctypes.cast(logits_out_address, P_FLOAT)
+            ctypes.cast(logits_out_address, P_FLOAT),
-        ), 'rwkv_eval_sequence_in_chunks failed, check stderr'
+        ):
+            raise ValueError("rwkv_eval_sequence_in_chunks failed, check stderr")
+
+    def rwkv_get_arch_version_major(self, ctx: RWKVContext) -> int:
+        """
+        Returns the major version used by the given model.
+
+        Parameters
+        ----------
+        ctx : RWKVContext
+            RWKV context obtained from rwkv_init_from_file.
+        """
+
+        return self.library.rwkv_get_arch_version_major(ctx.ptr)
+
+    def rwkv_get_arch_version_minor(self, ctx: RWKVContext) -> int:
+        """
+        Returns the minor version used by the given model.
+
+        Parameters
+        ----------
+        ctx : RWKVContext
+            RWKV context obtained from rwkv_init_from_file.
+        """
+
+        return self.library.rwkv_get_arch_version_minor(ctx.ptr)
 
     def rwkv_get_n_vocab(self, ctx: RWKVContext) -> int:
         """
@@ -358,7 +407,9 @@ class RWKVSharedLibrary:
 
         ctx.ptr = self.nullptr
 
-    def rwkv_quantize_model_file(self, model_file_path_in: str, model_file_path_out: str, format_name: str) -> None:
+    def rwkv_quantize_model_file(
+        self, model_file_path_in: str, model_file_path_out: str, format_name: str
+    ) -> None:
         """
         Quantizes FP32 or FP16 model to one of INT4 formats.
         Throws an exception in case of any error. Error messages would be printed to stderr.
@@ -373,20 +424,25 @@ class RWKVSharedLibrary:
             One of QUANTIZED_FORMAT_NAMES.
         """
 
-        assert format_name in QUANTIZED_FORMAT_NAMES, f'Unknown format name {format_name}, use one of {QUANTIZED_FORMAT_NAMES}'
+        if format_name not in QUANTIZED_FORMAT_NAMES:
+            raise ValueError(
+                f"Unknown format name {format_name}, use one of {QUANTIZED_FORMAT_NAMES}"
+            )
 
-        assert self.library.rwkv_quantize_model_file(
+        if not self.library.rwkv_quantize_model_file(
-            model_file_path_in.encode('utf-8'),
+            model_file_path_in.encode("utf-8"),
-            model_file_path_out.encode('utf-8'),
+            model_file_path_out.encode("utf-8"),
-            format_name.encode('utf-8')
+            format_name.encode("utf-8"),
-        ), 'rwkv_quantize_model_file failed, check stderr'
+        ):
+            raise ValueError("rwkv_quantize_model_file failed, check stderr")
 
     def rwkv_get_system_info_string(self) -> str:
         """
         Returns system information string.
         """
 
-        return self.library.rwkv_get_system_info_string().decode('utf-8')
+        return self.library.rwkv_get_system_info_string().decode("utf-8")
+
 
 def load_rwkv_shared_library() -> RWKVSharedLibrary:
     """
@@ -396,27 +452,27 @@ def load_rwkv_shared_library() -> RWKVSharedLibrary:
 
     file_name: str
 
-    if 'win32' in sys.platform or 'cygwin' in sys.platform:
+    if "win32" in sys.platform or "cygwin" in sys.platform:
-        file_name = 'rwkv.dll'
+        file_name = "rwkv.dll"
-    elif 'darwin' in sys.platform:
+    elif "darwin" in sys.platform:
-        file_name = 'librwkv.dylib'
+        file_name = "librwkv.dylib"
     else:
-        file_name = 'librwkv.so'
+        file_name = "librwkv.so"
 
     # Possible sub-paths to the library relative to the repo dir.
     child_paths: List[Callable[[pathlib.Path], pathlib.Path]] = [
         # No lookup for Debug config here.
         # I assume that if a user wants to debug the library,
         # they will be able to find the library and set the exact path explicitly.
-        lambda p: p / 'backend-python' / 'rwkv_pip' / 'cpp' / file_name,
+        lambda p: p / "backend-python" / "rwkv_pip" / "cpp" / file_name,
-        lambda p: p / 'bin' / 'Release' / file_name,
+        lambda p: p / "bin" / "Release" / file_name,
-        lambda p: p / 'bin' / file_name,
+        lambda p: p / "bin" / file_name,
         # Some people prefer to build in the "build" subdirectory.
-        lambda p: p / 'build' / 'bin' / 'Release' / file_name,
+        lambda p: p / "build" / "bin" / "Release" / file_name,
-        lambda p: p / 'build' / 'bin' / file_name,
+        lambda p: p / "build" / "bin" / file_name,
-        lambda p: p / 'build' / file_name,
+        lambda p: p / "build" / file_name,
         # Fallback.
-        lambda p: p / file_name
+        lambda p: p / file_name,
     ]
 
     working_dir: pathlib.Path = pathlib.Path(os.path.abspath(os.getcwd()))
@@ -430,7 +486,7 @@ def load_rwkv_shared_library() -> RWKVSharedLibrary:
         # .
         working_dir,
         # Repo dir relative to this Python file.
-        pathlib.Path(os.path.abspath(__file__)).parent.parent.parent
+        pathlib.Path(os.path.abspath(__file__)).parent.parent.parent,
     ]
 
     for parent_path in parent_paths:
@@ -440,5 +496,7 @@ def load_rwkv_shared_library() -> RWKVSharedLibrary:
             if os.path.isfile(full_path):
                 return RWKVSharedLibrary(str(full_path))
 
-    assert False, (f'Failed to find {file_name} automatically; '
+    raise ValueError(
-                   f'you need to find the library and create RWKVSharedLibrary specifying the path to it')
+        f"Failed to find {file_name} automatically; "
+        f"you need to find the library and create RWKVSharedLibrary specifying the path to it"
+    )

backend-python/rwkv_pip/webgpu/model.py

@@ -18,6 +18,7 @@ class RWKV:
         self.w["emb.weight"] = [0] * self.info.num_vocab
         self.version = str(self.info.version).lower()
         self.wrp = getattr(wrp, self.version)
+        self.version = float(self.version.replace("v", ""))
 
         layer = (
             int(s.lstrip("layer"))

backend-python/utils/rwkv.py

@@ -26,6 +26,7 @@ class AbstractRWKV(ABC):
         self.EOS_ID = 0
 
         self.name = "rwkv"
+        self.version = 4
         self.model = model
         self.pipeline = pipeline
         self.model_state = None
@@ -665,6 +666,7 @@ def RWKV(model: str, strategy: str, tokenizer: Union[str, None]) -> AbstractRWKV
         else:
             rwkv = TextRWKV(model, pipeline)
     rwkv.name = filename
+    rwkv.version = model.version
 
     return rwkv
 
@@ -677,7 +679,10 @@ class ModelConfigBody(BaseModel):
     frequency_penalty: float = Field(default=None, ge=-2, le=2)
     penalty_decay: float = Field(default=None, ge=0.99, le=0.999)
     top_k: int = Field(default=None, ge=0, le=25)
-    global_penalty: bool = Field(default=None)
+    global_penalty: bool = Field(
+        default=None,
+        description="When generating a response, whether to include the submitted prompt as a penalty factor. By turning this off, you will get the same generated results as official RWKV Gradio. If you find duplicate results in the generated results, turning this on can help avoid generating duplicates.",
+    )
 
     model_config = {
         "json_schema_extra": {

finetune/get_layer_and_embd.py

@@ -52,9 +52,13 @@ for x in keys:
     if "time_maa" in x:
         version = max(6, version)
 
+params = f"--vocab_size {vocab_size} --n_layer {n_layer} --n_embd {n_embd}"
+
 if version <= expected_max_version:
+    if version == 6:
+        params += ' --my_testing "x060"'
     print(
-        f"v{int(version)}/train.py --vocab_size {vocab_size} --n_layer {n_layer} --n_embd {n_embd}",
+        f"v{int(version)}/train.py {params}",
         end="",
     )
 else:

finetune/install-wsl-dep-and-train.sh

@@ -1,7 +1,7 @@
 echo $@
 
 if [[ ${cnMirror} == 1 ]]; then
-  export PIP_INDEX_URL="https://pypi.tuna.tsinghua.edu.cn/simple"
+  export PIP_INDEX_URL="https://mirrors.aliyun.com/pypi/simple"
   if grep -q "mirrors.aliyun.com" /etc/apt/sources.list; then
     echo "apt cnMirror already set"
   else
@@ -53,7 +53,7 @@ else
 fi
 
 echo "loading $loadModel"
-modelInfo=$(python3 ./finetune/get_layer_and_embd.py $loadModel 5.2)
+modelInfo=$(python3 ./finetune/get_layer_and_embd.py $loadModel 6.0)
 echo $modelInfo
 if [[ $modelInfo =~ "--n_layer" ]]; then
   sudo rm -rf /root/.cache/torch_extensions

finetune/lora/v6/cuda/wkv5_cuda.cu

@@ -0,0 +1,202 @@
+#include <stdio.h>
+#include <assert.h>
+#include "ATen/ATen.h"
+typedef at::BFloat16 bf16;
+
+template <typename F>
+__global__ void kernel_forward(const int B, const int T, const int C, const int H,
+                               const F *__restrict__ const _r, const F *__restrict__ const _k, const F *__restrict__ const _v, const float *__restrict__ _w, const F *__restrict__ _u,
+                               F *__restrict__ const _y)
+{
+    const int b = blockIdx.x / H;
+    const int h = blockIdx.x % H;
+    const int i = threadIdx.x;
+    _w += h*_N_;
+    _u += h*_N_;
+
+    __shared__ float r[_N_], k[_N_], u[_N_], w[_N_];
+    float state[_N_] = {0};
+
+    __syncthreads();
+    w[i] = _w[i];
+    u[i] = float(_u[i]);
+    __syncthreads();
+
+    for (int t = b*T*C + h*_N_ + i; t < (b+1)*T*C + h*_N_ + i; t += C)
+    {
+        __syncthreads();
+        r[i] = float(_r[t]);
+        k[i] = float(_k[t]);
+        __syncthreads();
+
+        const float v = float(_v[t]);
+        float y = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j+=4)
+        {
+            const float4& r_ = (float4&)(r[j]);
+            const float4& k_ = (float4&)(k[j]);
+            const float4& w_ = (float4&)(w[j]);
+            const float4& u_ = (float4&)(u[j]);
+            float4& s = (float4&)(state[j]);
+            float4 x;
+
+            x.x = k_.x * v;
+            x.y = k_.y * v;
+            x.z = k_.z * v;
+            x.w = k_.w * v;
+
+            y += r_.x * (u_.x * x.x + s.x);
+            y += r_.y * (u_.y * x.y + s.y);
+            y += r_.z * (u_.z * x.z + s.z);
+            y += r_.w * (u_.w * x.w + s.w);
+
+            s.x = s.x * w_.x + x.x;
+            s.y = s.y * w_.y + x.y;
+            s.z = s.z * w_.z + x.z;
+            s.w = s.w * w_.w + x.w;
+        }
+        _y[t] = F(y);
+    }
+}
+
+template <typename F>
+__global__ void kernel_backward(const int B, const int T, const int C, const int H,
+    const F *__restrict__ const _r, const F *__restrict__ const _k, const F *__restrict__ const _v, const float *__restrict__ _w, const float *__restrict__ __w, const F *__restrict__ _u, const F *__restrict__ const _gy,
+    F *__restrict__ const _gr, F *__restrict__ const _gk, F *__restrict__ const _gv, F *__restrict__ const _gw, F *__restrict__ const _gu)
+{
+    const int b = blockIdx.x / H;
+    const int h = blockIdx.x % H;
+    const int i = threadIdx.x;
+    _w += h*_N_;
+    _u += h*_N_;
+    __w += h*_N_;
+
+    __shared__ float w_[_N_], u_[_N_];
+    __shared__ float r[_N_], k[_N_], v[_N_], gy[_N_];
+    __syncthreads();
+    w_[i] = _w[i];
+    u_[i] = float(_u[i]);
+    __syncthreads();
+
+    const float w = w_[i];
+    const float ww = __w[i];
+    const float u = u_[i];
+
+    float state[_N_] = {0}, saaaa[_N_] = {0}, sbbbb[_N_] = {0}, scccc[_N_] = {0}, sdddd[_N_] = {0};
+
+    float gw = 0, gu = 0;
+    const int t000 = b*T*C + h*_N_ + i;
+    const int t111 = (b+1)*T*C + h*_N_ + i;
+    const int t222 = t111 - 2*C;
+
+    for (int t = t000; t < t111; t += C)
+    {
+        __syncthreads();
+        v[i] = float(_v[t]);
+        gy[i] = float(_gy[t]);
+        __syncthreads();
+
+        const float k = float(_k[t]);
+        float gr = 0, gu_ = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = state[j];
+            float x = k * v[j];
+
+            gr += (u * x + s) * gy[j];
+            gu_ += x * gy[j];
+            s = s * w + x;
+        }
+        _gr[t] = F(gr);
+        gu += float(_r[t]) * gu_;
+    }
+    _gu[b*C + h*_N_ + i] = F(gu);
+    
+    for (int t = t000; t < t222; t += C)
+    {
+        __syncthreads();
+        v[i] = float(_v[t]);
+        gy[i] = float(_gy[t + 2*C]);
+        __syncthreads();
+
+        const float k = float(_k[t]);
+        float gw_ = 0;
+        
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = saaaa[j];
+            float& s2 = sbbbb[j];
+            float x = k * v[j];
+            
+            float tmp = w * (x + s);
+            s = tmp;
+            s2 = tmp + w * s2;
+            gw_ += s2 * gy[j];
+        }
+        gw += float(_r[t + 2*C]) * gw_;
+    }    
+    _gw[b*C + h*_N_ + i] = F(ww * gw);
+
+    for (int t = t111 - C; t >= t000; t -= C)
+    {
+        __syncthreads();
+        v[i] = float(_v[t]);
+        gy[i] = float(_gy[t]);
+        __syncthreads();
+
+        const float rr = float(_r[t]);
+        float gk = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = scccc[j];
+            float x = rr * gy[j];
+            
+            gk += (u * x + s) * v[j];
+            s = x + s * w;
+        }
+        _gk[t] = F(gk);
+    }
+
+    for (int t = t111 - C; t >= t000; t -= C)
+    {
+        __syncthreads();
+        r[i] = float(_r[t]);
+        k[i] = float(_k[t]);
+        __syncthreads();
+
+        const float gyy = float(_gy[t]);
+        float gv = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = sdddd[j];
+            float x = gyy * r[j];
+            
+            gv += (u_[j] * x + s) * k[j];
+            s = x + s * w_[j];
+        }
+        _gv[t] = F(gv);
+    }
+}
+
+void cuda_forward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, float *w, bf16 *u, bf16 *y)
+{
+    assert(H*_N_ == C);
+    assert(_N_%4 == 0);
+    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, k, v, w, u, y);
+}
+
+void cuda_backward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, float *w, float *ww, bf16 *u, bf16 *gy, bf16 *gr, bf16 *gk, bf16 *gv, bf16 *gw, bf16 *gu)
+{
+    assert(H*_N_ == C);
+    assert(_N_%4 == 0);
+    kernel_backward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, k, v, w, ww, u, gy, gr, gk, gv, gw, gu);
+}

finetune/lora/v6/cuda/wkv5_op.cpp

@@ -0,0 +1,22 @@
+#include <torch/extension.h>
+#include "ATen/ATen.h"
+typedef at::BFloat16 bf16;
+
+void cuda_forward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, float *w, bf16 *u, bf16 *y);
+void cuda_backward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, float *w, float *ww, bf16 *u, bf16 *gy, bf16 *gr, bf16 *gk, bf16 *gv, bf16 *gw, bf16 *gu);
+
+void forward(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &y) {
+    cuda_forward(B, T, C, H, r.data_ptr<bf16>(), k.data_ptr<bf16>(), v.data_ptr<bf16>(), w.data_ptr<float>(), u.data_ptr<bf16>(), y.data_ptr<bf16>());
+}
+void backward(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &ww, torch::Tensor &u, torch::Tensor &gy, torch::Tensor &gr, torch::Tensor &gk, torch::Tensor &gv, torch::Tensor &gw, torch::Tensor &gu) {
+    cuda_backward(B, T, C, H, r.data_ptr<bf16>(), k.data_ptr<bf16>(), v.data_ptr<bf16>(), w.data_ptr<float>(), ww.data_ptr<float>(), u.data_ptr<bf16>(), gy.data_ptr<bf16>(), gr.data_ptr<bf16>(), gk.data_ptr<bf16>(), gv.data_ptr<bf16>(), gw.data_ptr<bf16>(), gu.data_ptr<bf16>());
+}
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("forward", &forward, "wkv5 forward");
+    m.def("backward", &backward, "wkv5 backward");
+}
+
+TORCH_LIBRARY(wkv5, m) {
+    m.def("forward", forward);
+    m.def("backward", backward);
+}

finetune/lora/v6/cuda/wkv6_cuda.cu

@@ -0,0 +1,242 @@
+#include <stdio.h>
+#include <assert.h>
+#include "ATen/ATen.h"
+typedef at::BFloat16 bf16;
+
+template <typename F>
+__global__ void kernel_forward(const int B, const int T, const int C, const int H,
+                               const F *__restrict__ const _r, const F *__restrict__ const _k, const F *__restrict__ const _v, const float *__restrict__ _w, const F *__restrict__ _u,
+                               F *__restrict__ const _y)
+{
+    const int b = blockIdx.x / H;
+    const int h = blockIdx.x % H;
+    const int i = threadIdx.x;
+    _u += h*_N_;
+
+    __shared__ float r[_N_], k[_N_], u[_N_], w[_N_];
+    float state[_N_] = {0};
+
+    __syncthreads();
+    u[i] = float(_u[i]);
+    __syncthreads();
+
+    for (int t = b*T*C + h*_N_ + i; t < (b+1)*T*C + h*_N_ + i; t += C)
+    {
+        __syncthreads();
+        w[i] = exp(_w[t]);
+        r[i] = float(_r[t]);
+        k[i] = float(_k[t]);
+        __syncthreads();
+
+        const float v = float(_v[t]);
+        float y = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j+=4)
+        {
+            const float4& r_ = (float4&)(r[j]);
+            const float4& k_ = (float4&)(k[j]);
+            const float4& w_ = (float4&)(w[j]);
+            const float4& u_ = (float4&)(u[j]);
+            float4& s = (float4&)(state[j]);
+            float4 x;
+
+            x.x = k_.x * v;
+            x.y = k_.y * v;
+            x.z = k_.z * v;
+            x.w = k_.w * v;
+
+            y += r_.x * (u_.x * x.x + s.x);
+            y += r_.y * (u_.y * x.y + s.y);
+            y += r_.z * (u_.z * x.z + s.z);
+            y += r_.w * (u_.w * x.w + s.w);
+
+            s.x = s.x * w_.x + x.x;
+            s.y = s.y * w_.y + x.y;
+            s.z = s.z * w_.z + x.z;
+            s.w = s.w * w_.w + x.w;
+        }
+        _y[t] = F(y);
+    }
+}
+
+template <typename F>
+__global__ void kernel_backward_111(const int B, const int T, const int C, const int H,
+    const F *__restrict__ const _r, const F *__restrict__ const _k, const F *__restrict__ const _v, const float *__restrict__ _w, const F *__restrict__ _u, const F *__restrict__ const _gy,
+    F *__restrict__ const _gr, F *__restrict__ const _gk, F *__restrict__ const _gv, F *__restrict__ const _gu)
+{
+    const int b = blockIdx.x / H;
+    const int h = blockIdx.x % H;
+    const int i = threadIdx.x;
+    _u += h*_N_;
+
+    __shared__ float u_[_N_];
+    __shared__ float r[_N_], k[_N_], v[_N_], w_[_N_], gy[_N_];
+    __syncthreads();
+    u_[i] = float(_u[i]);
+    __syncthreads();
+
+    const float u = u_[i];
+
+    float state[_N_] = {0}, scccc[_N_] = {0}, sdddd[_N_] = {0};
+
+    const int t_0 = b*T*C + h*_N_ + i;
+    const int t_T_1 = t_0 + (T-1)*C;
+    const int t_T = t_0 + T*C;
+
+    float gu = 0;
+    for (int t = t_0; t < t_T; t += C)
+    {
+        __syncthreads();
+        v[i] = float(_v[t]);
+        gy[i] = float(_gy[t]);
+        __syncthreads();
+
+        const float k = float(_k[t]);
+        const float w = exp(_w[t]);
+        float gr = 0, gu_ = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = state[j];
+            float x = k * v[j];
+
+            gr += (u * x + s) * gy[j];
+            gu_ += x * gy[j];
+            s = s * w + x;
+        }
+        _gr[t] = F(gr);
+        gu += float(_r[t]) * gu_;
+    }
+    _gu[b*C + h*_N_ + i] = F(gu);
+
+    for (int t = t_T_1; t >= t_0; t -= C)
+    {
+        __syncthreads();
+        v[i] = float(_v[t]);
+        gy[i] = float(_gy[t]);
+        __syncthreads();
+
+        const float rr = float(_r[t]);
+        const float w = exp(_w[t]);
+        float gk = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = scccc[j];
+            float x = rr * gy[j];
+            
+            gk += (u * x + s) * v[j];
+            s = x + s * w;
+        }
+        _gk[t] = F(gk);
+    }
+
+    for (int t = t_T_1; t >= t_0; t -= C)
+    {
+        __syncthreads();
+        r[i] = float(_r[t]);
+        k[i] = float(_k[t]);
+        w_[i] = exp(_w[t]);
+        __syncthreads();
+
+        const float gyy = float(_gy[t]);
+        float gv = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = sdddd[j];
+            float x = gyy * r[j];
+            
+            gv += (u_[j] * x + s) * k[j];
+            s = x + s * w_[j];
+        }
+        _gv[t] = F(gv);
+    }
+}
+
+template <typename F>
+__global__ void kernel_backward_222(const int B, const int T, const int C, const int H,
+    const F *__restrict__ const _r, const F *__restrict__ const _k, const F *__restrict__ const _v, const float *__restrict__ _w, const F *__restrict__ _u, const F *__restrict__ const _gy,
+    F *__restrict__ const _gw)
+{
+    const int b = blockIdx.x / H;
+    const int h = blockIdx.x % H;
+    const int i = threadIdx.x;
+
+    __shared__ float v[_N_], gy[_N_];
+    float saaaa[_N_] = {0}, sbbbb[_T_-2] = {0}, scccc[_N_] = {0};
+
+    const int t_0 = b*T*C + h*_N_ + i;
+    const int t_1 = t_0 + C;
+    const int t_2 = t_0 + 2*C;
+    const int t_T_1 = t_0 + (T-1)*C;
+
+    for (int t = t_T_1; t > t_1; t -= C)
+    {
+        __syncthreads();
+        gy[i] = float(_gy[t]);
+        v[i] = float(_v[t-2*C]);
+        __syncthreads();
+
+        const float r = float(_r[t]);
+        const float w = exp(_w[t-C]);
+        float sum = 0.0f;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = saaaa[j];
+            float x = r * gy[j];
+            s = (s + x) * w;
+            sum += s * v[j];
+        }
+        sbbbb[(t-t_2)/C] = sum * float(_k[t-2*C]);
+    }
+
+    float sss = sbbbb[0];
+    _gw[t_0] = 0;
+    _gw[t_1] = F(sss * _w[t_1]);
+
+    for (int t = t_2; t < t_T_1; t += C)
+    {
+        __syncthreads();
+        gy[i] = float(_gy[t]);
+        v[i] = float(_v[t-2*C]);
+        __syncthreads();
+
+        const float w = exp(_w[t-C]);
+        const float k = float(_k[t-2*C]);
+        float sum = 0.0f;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = scccc[j];
+            float x = k * v[j];
+            s = (s + x) * w;
+            sum += s * gy[j];
+        }
+        sss += sbbbb[(t-t_1)/C] - (sum * float(_r[t]));
+        _gw[t] = F(sss * _w[t]);
+    }
+    _gw[t_T_1] = 0;
+}
+
+void cuda_forward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, float *w, bf16 *u, bf16 *y)
+{
+    assert(H*_N_ == C);
+    assert(_N_%4 == 0);
+    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, k, v, w, u, y);
+}
+
+void cuda_backward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, float *w, bf16 *u, bf16 *gy, bf16 *gr, bf16 *gk, bf16 *gv, bf16 *gw, bf16 *gu)
+{
+    assert(H*_N_ == C);
+    assert(_N_%4 == 0);
+    kernel_backward_111<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, k, v, w, u, gy, gr, gk, gv, gu);
+    kernel_backward_222<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, k, v, w, u, gy, gw);
+}

finetune/lora/v6/cuda/wkv6_op.cpp

@@ -0,0 +1,22 @@
+#include <torch/extension.h>
+#include "ATen/ATen.h"
+typedef at::BFloat16 bf16;
+
+void cuda_forward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, float *w, bf16 *u, bf16 *y);
+void cuda_backward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, float *w, bf16 *u, bf16 *gy, bf16 *gr, bf16 *gk, bf16 *gv, bf16 *gw, bf16 *gu);
+
+void forward(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &y) {
+    cuda_forward(B, T, C, H, r.data_ptr<bf16>(), k.data_ptr<bf16>(), v.data_ptr<bf16>(), w.data_ptr<float>(), u.data_ptr<bf16>(), y.data_ptr<bf16>());
+}
+void backward(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &gy, torch::Tensor &gr, torch::Tensor &gk, torch::Tensor &gv, torch::Tensor &gw, torch::Tensor &gu) {
+    cuda_backward(B, T, C, H, r.data_ptr<bf16>(), k.data_ptr<bf16>(), v.data_ptr<bf16>(), w.data_ptr<float>(), u.data_ptr<bf16>(), gy.data_ptr<bf16>(), gr.data_ptr<bf16>(), gk.data_ptr<bf16>(), gv.data_ptr<bf16>(), gw.data_ptr<bf16>(), gu.data_ptr<bf16>());
+}
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("forward", &forward, "wkv6 forward");
+    m.def("backward", &backward, "wkv6 backward");
+}
+
+TORCH_LIBRARY(wkv6, m) {
+    m.def("forward", forward);
+    m.def("backward", backward);
+}

finetune/lora/v6/src/binidx.py

@@ -0,0 +1,303 @@
+from lib2to3.pgen2 import token
+import os
+import torch
+import numpy as np
+import shutil
+import struct
+from functools import lru_cache
+from itertools import accumulate
+
+
+def print_rank_0(*message):
+    pass
+    # """If distributed is initialized print only on rank 0."""
+    # if torch.distributed.is_initialized():
+    #     if torch.distributed.get_rank() == 0:
+    #         print(*message, flush=True)
+    # else:
+    #     print(*message, flush=True)
+
+
+def _warmup_mmap_file(path):
+    pass
+    # with open(path, "rb") as stream:
+    #     while stream.read(100 * 1024 * 1024):
+    #         pass
+
+
+dtypes = {
+    1: np.uint8,
+    2: np.int8,
+    3: np.int16,
+    4: np.int32,
+    5: np.int64,
+    6: float,
+    7: np.double,
+    8: np.uint16,
+}
+
+
+def code(dtype):
+    for k in dtypes.keys():
+        if dtypes[k] == dtype:
+            return k
+    raise ValueError(dtype)
+
+
+def index_file_path(prefix_path):
+    return prefix_path + ".idx"
+
+
+def data_file_path(prefix_path):
+    return prefix_path + ".bin"
+
+
+class MMapIndexedDataset(torch.utils.data.Dataset):
+    class Index(object):
+        _HDR_MAGIC = b"MMIDIDX\x00\x00"
+
+        @classmethod
+        def writer(cls, path, dtype):
+            class _Writer(object):
+                def __enter__(self):
+                    self._file = open(path, "wb")
+
+                    # Write Magic string so we can check the file format then opening it again.
+                    self._file.write(cls._HDR_MAGIC)
+                    # Write version number
+                    # Little endian unsigned 64 Bit integer
+                    self._file.write(struct.pack("<Q", 1))
+                    # Little endian unsigned 8 Bit integer
+                    self._file.write(struct.pack("<B", code(dtype)))
+
+                    return self
+
+                @staticmethod
+                def _get_pointers(sizes):
+                    dtype_size = dtype().itemsize
+                    address = 0
+                    pointers = []
+
+                    for size in sizes:
+                        pointers.append(address)
+                        address += size * dtype_size
+
+                    return pointers
+
+                def write(self, sizes, doc_idx):
+                    pointers = self._get_pointers(sizes)
+
+                    # Little endian unsigned 64 Bit integer
+                    self._file.write(struct.pack("<Q", len(sizes)))
+                    # Little endian unsigned 64 Bit integer
+                    self._file.write(struct.pack("<Q", len(doc_idx)))
+
+                    sizes = np.array(sizes, dtype=np.int32)
+                    self._file.write(sizes.tobytes(order="C"))
+                    del sizes
+
+                    pointers = np.array(pointers, dtype=np.int64)
+                    self._file.write(pointers.tobytes(order="C"))
+                    del pointers
+
+                    doc_idx = np.array(doc_idx, dtype=np.int64)
+                    self._file.write(doc_idx.tobytes(order="C"))
+
+                def __exit__(self, exc_type, exc_val, exc_tb):
+                    self._file.close()
+
+            return _Writer()
+
+        def __init__(self, path, skip_warmup=False):
+            with open(path, "rb") as stream:
+                magic_test = stream.read(9)
+                assert self._HDR_MAGIC == magic_test, (
+                    "Index file doesn't match expected format. "
+                    "Make sure that --dataset-impl is configured properly."
+                )
+                # Little endian unsigned 64 Bit integer
+                version = struct.unpack("<Q", stream.read(8))
+                assert (1,) == version
+
+                # Little endian unsigned 8 Bit integer
+                (dtype_code,) = struct.unpack("<B", stream.read(1))
+                self._dtype = dtypes[dtype_code]
+                self._dtype_size = self._dtype().itemsize
+
+                self._len = struct.unpack("<Q", stream.read(8))[0]
+                self._doc_count = struct.unpack("<Q", stream.read(8))[0]
+                offset = stream.tell()
+
+            if not skip_warmup:
+                print_rank_0("    warming up index mmap file...")
+                _warmup_mmap_file(path)
+
+            self._bin_buffer_mmap = np.memmap(path, mode="r", order="C")
+            self._bin_buffer = memoryview(self._bin_buffer_mmap)
+            print_rank_0("    reading sizes...")
+            self._sizes = np.frombuffer(
+                self._bin_buffer, dtype=np.int32, count=self._len, offset=offset
+            )
+            print_rank_0("    reading pointers...")
+            self._pointers = np.frombuffer(
+                self._bin_buffer,
+                dtype=np.int64,
+                count=self._len,
+                offset=offset + self._sizes.nbytes,
+            )
+            print_rank_0("    reading document index...")
+            self._doc_idx = np.frombuffer(
+                self._bin_buffer,
+                dtype=np.int64,
+                count=self._doc_count,
+                offset=offset + self._sizes.nbytes + self._pointers.nbytes,
+            )
+
+        def __del__(self):
+            self._bin_buffer_mmap._mmap.close()
+            del self._bin_buffer_mmap
+
+        @property
+        def dtype(self):
+            return self._dtype
+
+        @property
+        def sizes(self):
+            return self._sizes
+
+        @property
+        def doc_idx(self):
+            return self._doc_idx
+
+        @lru_cache(maxsize=8)
+        def __getitem__(self, i):
+            return self._pointers[i], self._sizes[i]
+
+        def __len__(self):
+            return self._len
+
+    def __init__(self, path, skip_warmup=False):
+        super().__init__()
+
+        self._path = None
+        self._index = None
+        self._bin_buffer = None
+
+        self._do_init(path, skip_warmup)
+
+    def __getstate__(self):
+        return self._path
+
+    def __setstate__(self, state):
+        self._do_init(state)
+
+    def _do_init(self, path, skip_warmup):
+        self._path = path
+        self._index = self.Index(index_file_path(self._path), skip_warmup)
+
+        if not skip_warmup:
+            print_rank_0("    warming up data mmap file...")
+            _warmup_mmap_file(data_file_path(self._path))
+        print_rank_0("    creating numpy buffer of mmap...")
+        self._bin_buffer_mmap = np.memmap(
+            data_file_path(self._path), mode="r", order="C"
+        )
+        print_rank_0("    creating memory view of numpy buffer...")
+        self._bin_buffer = memoryview(self._bin_buffer_mmap)
+
+    def __del__(self):
+        self._bin_buffer_mmap._mmap.close()
+        del self._bin_buffer_mmap
+        del self._index
+
+    def __len__(self):
+        return len(self._index)
+
+    # @lru_cache(maxsize=8)
+    def __getitem__(self, idx):
+        if isinstance(idx, int):
+            ptr, size = self._index[idx]
+            np_array = np.frombuffer(
+                self._bin_buffer, dtype=self._index.dtype, count=size, offset=ptr
+            )
+            return np_array
+        elif isinstance(idx, slice):
+            start, stop, step = idx.indices(len(self))
+            if step != 1:
+                raise ValueError("Slices into indexed_dataset must be contiguous")
+            ptr = self._index._pointers[start]
+            sizes = self._index._sizes[idx]
+            offsets = list(accumulate(sizes))
+            total_size = sum(sizes)
+            np_array = np.frombuffer(
+                self._bin_buffer, dtype=self._index.dtype, count=total_size, offset=ptr
+            )
+            sents = np.split(np_array, offsets[:-1])
+            return sents
+
+    def get(self, idx, offset=0, length=None):
+        """Retrieves a single item from the dataset with the option to only
+        return a portion of the item.
+
+        get(idx) is the same as [idx] but get() does not support slicing.
+        """
+        ptr, size = self._index[idx]
+        if length is None:
+            length = size - offset
+        ptr += offset * np.dtype(self._index.dtype).itemsize
+        np_array = np.frombuffer(
+            self._bin_buffer, dtype=self._index.dtype, count=length, offset=ptr
+        )
+        return np_array
+
+    def pad(self, idx, length=None):
+        ptr, size = self._index[idx]
+        try:
+            np_array = np.frombuffer(
+                self._bin_buffer, dtype=self._index.dtype, count=length, offset=ptr
+            )
+        except:
+            np_array = np.frombuffer(
+                self._bin_buffer, dtype=self._index.dtype, count=size, offset=ptr
+            )
+            ptr0, _ = self._index[0]
+            np_array0 = np.frombuffer(
+                self._bin_buffer,
+                dtype=self._index.dtype,
+                count=length - size,
+                offset=ptr0,
+            )
+            np_array = np.append(np_array, np_array0)
+        return np_array
+
+    def only(self, idx):
+        ptr, size = self._index[idx]
+        np_array = np.frombuffer(
+            self._bin_buffer, dtype=self._index.dtype, count=size, offset=ptr
+        )
+
+        return np_array
+
+    @property
+    def sizes(self):
+        return self._index.sizes
+
+    @property
+    def doc_idx(self):
+        return self._index.doc_idx
+
+    def get_doc_idx(self):
+        return self._index._doc_idx
+
+    def set_doc_idx(self, doc_idx_):
+        self._index._doc_idx = doc_idx_
+
+    @property
+    def supports_prefetch(self):
+        return False
+
+    @staticmethod
+    def exists(path):
+        return os.path.exists(index_file_path(path)) and os.path.exists(
+            data_file_path(path)
+        )

finetune/lora/v6/src/dataset.py

@@ -0,0 +1,242 @@
+########################################################################################################
+# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM
+########################################################################################################
+
+import json, math, random, os, sys
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from pytorch_lightning.utilities import rank_zero_info
+from .binidx import MMapIndexedDataset
+from .utils import MaybeIsPrime
+
+
+class MyDataset(Dataset):
+    def __init__(self, args):
+        self.args = args
+
+        if args.data_type == "binidx":
+            self.vocab_size = args.vocab_size
+            rank_zero_info(
+                f"Current vocab size = {self.vocab_size} (make sure it's correct)"
+            )
+
+            if args.my_pile_version == 1:
+                self.data = MMapIndexedDataset(args.data_file)
+                self.data_size = (
+                    len(self.data._bin_buffer) // self.data._index._dtype_size
+                )
+                rank_zero_info(f"Data has {self.data_size} tokens.")
+            elif args.my_pile_version == 2:
+                data_list = (
+                    open(args.data_file, "r", encoding="utf-8")
+                    .read()
+                    .strip()
+                    .split("\n")
+                )
+                data_list = [i.strip().split(" ") for i in data_list]
+                self.data = []
+                self.data_size = int(data_list[-1][-1])
+                rank_zero_info(f"Data has {self.data_size} chunks.")
+                for d in data_list:
+                    data = MMapIndexedDataset(d[0])
+                    data_size = len(data._bin_buffer) // data._index._dtype_size
+                    assert (data_size - args.ctx_len) == int(d[1])
+                    self.data += [[int(d[-1]), int(d[1]), data]]
+                # rank_zero_info(self.data)
+
+            if args.my_qa_mask > 0:
+                # self.data_pile = MMapIndexedDataset('/fsx/pile/pile_20B_tokenizer_text_document')
+                self.data_pile = MMapIndexedDataset(
+                    "/fsx/pile_deduped/pile_0.87_deduped_text_document"
+                )
+                self.data_pile_size = (
+                    len(self.data_pile._bin_buffer) // self.data._index._dtype_size
+                )
+            else:
+                self.data_pile = None
+                self.data_pile_size = 0
+
+            if args.my_pile_stage > 0:
+                # assert self.data_size == 332115325534 and self.vocab_size == 50277
+                self.samples_per_epoch = args.epoch_steps * args.real_bsz
+                assert self.samples_per_epoch == 40320
+                rank_zero_info(
+                    f"########## Pile 20b-tokenized stage {args.my_pile_stage} ##########"
+                )
+                dataset_slot = self.data_size // args.ctx_len
+                if args.my_pile_stage != 4:
+                    assert MaybeIsPrime(args.magic_prime)
+                    assert args.magic_prime % 3 == 2
+                    assert (
+                        args.magic_prime / dataset_slot > 0.99
+                        and args.magic_prime / dataset_slot <= 1
+                    )
+        elif args.data_type == "numpy":
+            self.data = np.load(args.data_file).astype("int")
+            self.vocab_size = args.vocab_size
+            rank_zero_info(
+                f"Current vocab size = {self.vocab_size} (make sure it's correct)"
+            )
+            self.data_size = len(self.data)
+            rank_zero_info(f"Data has {self.data_size} tokens.")
+        elif args.data_type == "uint16":
+            self.data = (
+                np.fromfile(args.data_file, dtype=np.uint16)
+                .astype("int32")
+                .reshape(-1, args.my_sample_len)
+            )
+            self.vocab_size = args.vocab_size
+            rank_zero_info(
+                f"Current vocab size = {self.vocab_size} (make sure it's correct)"
+            )
+            self.data_size = self.data.shape[0]
+            rank_zero_info(f"Data has {self.data_size} samples.")
+        else:
+            if args.data_type == "dummy":
+                rank_zero_info("Building dummy data...")
+                self.data = ""
+                for i in range(100000):
+                    aa = (i) % 10000
+                    bb = (i * i) % 10000
+                    cc = aa + bb
+                    self.data += f".{aa}+{bb}={cc}."
+            else:
+                self.data = open(args.data_file, "r", encoding=args.data_type).read()
+            rank_zero_info("Building token list...")
+            unique = sorted(list(set(self.data)))
+            self.vocab_size = len(unique)
+            # rank_zero_info()
+            # for u in unique:
+            #     print(u, end=' ')
+            # rank_zero_info('\n\n')
+            xx = 0
+            xxObj = {}
+            for u in unique:
+                xxObj[xx] = u
+                xx += 1
+            with open(
+                f"{args.proj_dir}/vocab.json", "w", encoding="utf-8"
+            ) as vocab_file:
+                vocab_file.write(json.dumps(xxObj, ensure_ascii=False))
+            self.data_size = len(self.data)
+            rank_zero_info(
+                f"Data has {self.data_size} tokens, {self.vocab_size} vocab size."
+            )
+            self.stoi = {ch: i for i, ch in enumerate(unique)}
+            self.itos = {i: ch for i, ch in enumerate(unique)}
+
+    def __len__(self):
+        return self.args.epoch_steps * self.args.micro_bsz
+
+    def __getitem__(self, idx):
+        args = self.args
+        rank = self.global_rank
+        epoch = self.real_epoch
+        world_size = self.world_size
+        # print(f"epoch {epoch} idx {idx} rank {rank}/{world_size}")
+
+        if args.data_type == "uint16":
+            i = np.random.randint(0, self.data_size - 1)
+            dix = self.data[i]
+            x = torch.tensor(dix[:-1], dtype=torch.long)
+            y = torch.tensor(dix[1:], dtype=torch.long)
+        else:
+            ctx_len = args.ctx_len
+            req_len = ctx_len + 1
+            magic_prime = args.magic_prime
+            data = self.data
+
+            if args.my_pile_stage > 0:
+                ii = 1 + epoch * self.samples_per_epoch + (idx * world_size) + rank
+
+                if args.my_qa_mask > 0:
+                    ii_orig = ii
+                    if ii % 2 == 0:
+                        ii = -1
+                        data = self.data_pile
+                    else:
+                        ii = ii // 2
+                if data == self.data_pile:
+                    i = np.random.randint(0, self.data_pile_size - req_len)
+                else:
+                    if args.my_pile_stage == 4 or ii < args.my_random_steps:
+                        # cheat: pick a random spot in dataset
+                        if args.my_pile_version == 1:
+                            i = np.random.randint(0, self.data_size - req_len)
+                        else:
+                            i = np.random.randint(0, self.data_size)
+                    else:
+                        ii = ii - args.my_random_steps
+                        factor = (math.sqrt(5) - 1) / 2
+                        factor = int(magic_prime * factor)
+                        i = ((factor * ii * ii * ii) % magic_prime) * ctx_len
+                        i = i + args.my_pile_shift
+                # print(f"epoch {epoch} idx {idx} rank {rank}/{world_size} ii {ii} pos {round(i / self.data_size, 3)}")
+            else:
+                # cheat: pick a random spot in dataset
+                i = np.random.randint(0, self.data_size - req_len)
+
+            if args.data_type == "binidx":
+                if args.my_pile_version == 1:
+                    dix = data.get(idx=0, offset=i, length=req_len).astype(int)
+                    # dix = data.pad(idx=idx, length=req_len).astype(int)
+                else:
+                    # self.data : cutoff, chunk_count, data
+                    for j in range(len(data)):
+                        if i < data[j][0]:
+                            ii = i
+                            i = (i - (data[j - 1][0] if j > 0 else 0)) % data[j][1]
+                            dix = (
+                                data[j][2]
+                                .get(idx=0, offset=i, length=req_len)
+                                .astype(int)
+                            )
+                            # print(ii, j, i)
+                            break
+            elif args.data_type == "numpy":
+                dix = data[i : i + req_len]
+            else:
+                dix = [self.stoi[s] for s in data[i : i + req_len]]
+
+            if args.my_qa_mask == 1:
+                if data == self.data_pile:
+                    z = [1] * ctx_len
+                else:
+                    z = [0] * ctx_len
+                    z_sum = 0
+                    isGood = False
+                    for i in range(3, ctx_len):
+                        if (
+                            dix[i] == 27
+                            and dix[i - 1] == 34
+                            and dix[i - 2] == 187
+                            and dix[i - 3] == 187
+                        ):
+                            isGood = True
+                        if dix[i] == 0:
+                            isGood = False
+                        if isGood:
+                            z[i] = 1
+                            z_sum += 1
+                    if z_sum == 0:
+                        z = [1] * ctx_len
+                        i = np.random.randint(0, self.data_pile_size - req_len)
+                        dix = self.data_pile.get(
+                            idx=0, offset=i, length=req_len
+                        ).astype(int)
+                z = torch.tensor(z, dtype=torch.bfloat16)
+
+            x = torch.tensor(dix[:-1], dtype=torch.long)
+            y = torch.tensor(dix[1:], dtype=torch.long)
+
+            # if ii_orig < 50:
+            #     # if rank == 1:
+            #     print('rank', rank, 'i', ii_orig, ii, i, 'x', x[:5], '...', x[-5:])
+            # else:
+            #     exit(0)
+
+            if args.my_qa_mask == 1:
+                return x, y, z
+
+            return x, y

finetune/lora/v6/src/trainer.py

@@ -0,0 +1,310 @@
+import os, math, time, datetime, subprocess
+import torch
+from torch.utils.data import DataLoader
+import pytorch_lightning as pl
+from pytorch_lightning.utilities import rank_zero_info, rank_zero_only
+from .model import LORA_CONFIG
+
+
+def my_save(args, trainer, dd, ff):
+    if "14b-run1" in ff:
+        fn = ff.split("/")[-1]
+        fff = "/dev/shm/" + fn
+        torch.save(dd, fff)
+        subprocess.Popen(f" aws s3 mv {fff} s3://rwkv-14b-4k/{fn} --quiet", shell=True)
+    elif ("world/14b" in ff) or ("world/7b" in ff):
+        aa = ff.split("/")[1]
+        fn = ff.split("/")[-1]
+        fff = f"/dev/shm/{aa}-{fn}"
+        torch.save(dd, fff)
+        subprocess.Popen(
+            f" aws s3 mv {fff} s3://rwkv-world/{aa}-{fn} --quiet", shell=True
+        )
+    else:
+        if "deepspeed_stage_3" in args.strategy:
+            trainer.save_checkpoint(ff, weights_only=True)
+        else:
+            torch.save(dd, ff)
+
+
+class train_callback(pl.Callback):
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+
+    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
+        args = self.args
+        # if args.cuda_cleanup > 0:
+        #     torch.cuda.empty_cache()
+        real_step = trainer.global_step + args.epoch_begin * args.epoch_steps
+
+        # LR schedule
+        w_step = args.warmup_steps
+        if args.lr_final == args.lr_init or args.epoch_count == 0:
+            lr = args.lr_init
+        else:
+            decay_step = real_step - args.my_pile_edecay * args.epoch_steps
+            decay_total = (args.epoch_count - args.my_pile_edecay) * args.epoch_steps
+            progress = (decay_step - w_step + 1) / (decay_total - w_step)
+            progress = min(1, max(0, progress))
+
+            if args.lr_final == 0 or args.lr_init == 0:  # linear decay
+                lr = args.lr_init + (args.lr_final - args.lr_init) * progress
+            else:  # exp decay
+                lr = args.lr_init * math.exp(
+                    math.log(args.lr_final / args.lr_init) * pow(progress, 1)
+                )
+            # if trainer.is_global_zero:
+            #     print(trainer.global_step, decay_step, decay_total, w_step, progress, lr)
+
+        if args.my_exit_tokens != 0:  # cosine decay
+            real_tokens = real_step * args.ctx_len * args.real_bsz
+            warmup_tokens = w_step * args.ctx_len * args.real_bsz
+            progress = (real_tokens - warmup_tokens) / (
+                abs(args.my_exit_tokens) - warmup_tokens
+            )
+            progress = max(0, min(1, progress))
+            lr_final_factor = args.lr_final / args.lr_init
+            lr_mult = (0.5 + lr_final_factor / 2) + (
+                0.5 - lr_final_factor / 2
+            ) * math.cos(math.pi * progress)
+            if args.my_exit_tokens > 0:
+                lr = args.lr_init * lr_mult
+            else:
+                lr = (lr + args.lr_init * lr_mult) / 2
+            if progress >= 1:
+                if (trainer.is_global_zero) or ("deepspeed_stage_3" in args.strategy):
+                    my_save(
+                        args,
+                        trainer,
+                        pl_module.state_dict(),
+                        f"{args.proj_dir}/rwkv-final.pth",
+                    )
+                    exit(0)
+        if trainer.global_step < w_step:
+            lr = lr * (0.2 + 0.8 * trainer.global_step / w_step)
+
+        if args.weight_decay_final > 0:
+            wd_now = args.weight_decay * math.exp(
+                math.log(args.weight_decay_final / args.weight_decay) * progress
+            )
+        else:
+            wd_now = args.weight_decay
+
+        for param_group in trainer.optimizers[0].param_groups:
+            if param_group["weight_decay"] > 0:
+                param_group["weight_decay"] = wd_now
+            if args.layerwise_lr > 0:
+                param_group["lr"] = lr * param_group["my_lr_scale"]
+                # print(param_group["lr"], param_group["my_lr_scale"])
+            else:
+                param_group["lr"] = lr
+
+        trainer.my_lr = lr
+        trainer.my_wd = wd_now
+        # rank_zero_info(f"{real_step} {lr}")
+
+        if trainer.global_step == 0:
+            if trainer.is_global_zero:  # logging
+                trainer.my_loss_sum = 0
+                trainer.my_loss_count = 0
+                trainer.my_log = open(args.proj_dir + "/train_log.txt", "a")
+                trainer.my_log.write(
+                    f"NEW RUN {args.my_timestamp}\n{vars(self.args)}\n"
+                )
+                try:
+                    print(f"\n{trainer.strategy.config}\n")
+                    trainer.my_log.write(f"{trainer.strategy.config}\n")
+                except:
+                    pass
+                trainer.my_log.flush()
+                if len(args.wandb) > 0:
+                    print("Login to wandb...")
+                    import wandb
+
+                    wandb.init(
+                        project=args.wandb,
+                        name=args.run_name + " " + args.my_timestamp,
+                        config=args,
+                        save_code=False,
+                    )
+                    trainer.my_wandb = wandb
+
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
+        args = self.args
+        token_per_step = args.ctx_len * args.real_bsz
+        real_step = trainer.global_step + args.epoch_begin * args.epoch_steps
+        if trainer.is_global_zero:  # logging
+            t_now = time.time_ns()
+            kt_s = 0
+            try:
+                t_cost = (t_now - trainer.my_time_ns) / 1e9
+                kt_s = token_per_step / t_cost / 1000
+                self.log("REAL it/s", 1.0 / t_cost, prog_bar=True, on_step=True)
+                self.log("Kt/s", kt_s, prog_bar=True, on_step=True)
+            except:
+                pass
+            trainer.my_time_ns = t_now
+            if pl.__version__[0] == "2":
+                trainer.my_loss = outputs["loss"]
+            else:
+                trainer.my_loss = trainer.my_loss_all.float().mean().item()
+            trainer.my_loss_sum += trainer.my_loss
+            trainer.my_loss_count += 1
+            trainer.my_epoch_loss = trainer.my_loss_sum / trainer.my_loss_count
+            self.log("lr", trainer.my_lr, prog_bar=True, on_step=True)
+            self.log("loss", trainer.my_epoch_loss, prog_bar=True, on_step=True)
+            # self.log("s", real_step, prog_bar=True, on_step=True)
+
+            if len(args.wandb) > 0:
+                lll = {
+                    "loss": trainer.my_loss,
+                    "lr": trainer.my_lr,
+                    "wd": trainer.my_wd,
+                    "Gtokens": real_step * token_per_step / 1e9,
+                }
+                if kt_s > 0:
+                    lll["kt/s"] = kt_s
+                trainer.my_wandb.log(lll, step=int(real_step))
+        if (trainer.is_global_zero) or (
+            "deepspeed_stage_3" in args.strategy
+        ):  # save pth
+            if args.magic_prime > 0:
+                expand_factor = 2 if args.my_qa_mask > 0 else 1
+                if int(real_step) == int(
+                    args.magic_prime * expand_factor // args.real_bsz
+                ) - 1 + int(args.my_random_steps):
+                    to_save_dict = pl_module.state_dict()
+                    my_save(
+                        args,
+                        trainer,
+                        to_save_dict,
+                        f"{args.proj_dir}/rwkv-final.pth",
+                    )
+        # if args.batch_save==batch_idx :
+        #     to_save_dict = pl_module.state_dict()
+        #     for name, state in to_save_dict.items():
+        #         if 'img' in name:
+        #             to_save_dict[name] = state
+        #     try:
+        #             my_save(
+        #                 args, trainer,
+        #                 to_save_dict,
+        #                 f"{args.proj_dir}/rwkv-{args.epoch_begin + trainer.current_epoch}-{batch_idx}.pth",
+        #             )
+        #     except Exception as e:
+        #         print('Error\n\n', e, '\n\n')
+
+    def on_train_epoch_start(self, trainer, pl_module):
+        args = self.args
+        if pl.__version__[0] == "2":
+            dataset = trainer.train_dataloader.dataset
+        else:
+            dataset = trainer.train_dataloader.dataset.datasets
+        assert "MyDataset" in str(dataset)
+        dataset.global_rank = trainer.global_rank
+        dataset.real_epoch = int(args.epoch_begin + trainer.current_epoch)
+        dataset.world_size = trainer.world_size
+        # print(f'########## world_size {dataset.world_size} global_rank {dataset.global_rank} real_epoch {dataset.real_epoch} ##########')
+
+    def on_train_epoch_end(self, trainer, pl_module):
+        args = self.args
+        to_save_dict = {}
+        if (trainer.is_global_zero) or (
+            "deepspeed_stage_3" in args.strategy
+        ):  # save pth
+            if (
+                args.epoch_save > 0 and trainer.current_epoch % args.epoch_save == 0
+            ) or (trainer.current_epoch == args.epoch_count - 1):
+                if args.data_type == "wds_img":
+                    raw_dict = pl_module.state_dict()
+                    for k in raw_dict:
+                        if k.startswith("encoder.") or k.startswith("decoder."):
+                            to_save_dict[k] = raw_dict[k]
+                else:
+                    to_save_dict = pl_module.state_dict()
+
+                if args.data_type == "img" and not args.lora:
+                    for name, state in to_save_dict.items():
+                        if "img" in name:
+                            to_save_dict[name] = state
+
+                if args.lora:
+                    enable_time_finetune = "time" in LORA_CONFIG["parts"]
+                    enable_ln_finetune = "ln" in LORA_CONFIG["parts"]
+                    lora_dict = {}
+                    for name, state in to_save_dict.items():
+                        if "img" in name:
+                            lora_dict[name] = state
+                        if (
+                            ".lora_" in name
+                            or (enable_time_finetune and ".time_" in name)
+                            or (enable_ln_finetune and ".ln" in name)
+                        ):
+                            lora_dict[name] = state
+                    to_save_dict = lora_dict
+
+                try:
+                    my_save(
+                        args,
+                        trainer,
+                        to_save_dict,
+                        f"{args.proj_dir}/rwkv-{args.epoch_begin + trainer.current_epoch}.pth",
+                    )
+                except Exception as e:
+                    print("Error\n\n", e, "\n\n")
+
+        if trainer.is_global_zero:  # logging
+            trainer.my_log.write(
+                f"{args.epoch_begin + trainer.current_epoch} {trainer.my_epoch_loss:.6f} {math.exp(trainer.my_epoch_loss):.4f} {trainer.my_lr:.8f} {datetime.datetime.now()} {trainer.current_epoch}\n"
+            )
+            trainer.my_log.flush()
+
+            trainer.my_loss_sum = 0
+            trainer.my_loss_count = 0
+            if (args.epoch_begin + trainer.current_epoch) >= args.my_exit:
+                exit(0)
+
+
+@rank_zero_only
+def generate_init_weight(model, init_weight_name):
+    mm = model.generate_init_weight()
+
+    if model.args.my_pile_stage == 1:
+        if len(model.args.load_model) > 0:
+            print(f"Combine weights from {model.args.load_model}...")
+            load_dict = torch.load(model.args.load_model, map_location="cpu")
+            for k in load_dict:
+                try:
+                    assert k in mm
+                except:
+                    print("missing", k)
+                    exit(0)
+                src = load_dict[k]
+                try:
+                    mm[k] = src.reshape(mm[k].shape)
+                except:
+                    tmp = mm[k].squeeze().clone()
+                    print(k, src.shape, "-->", mm[k].shape)
+                    ss = src.shape[0]
+                    dd = tmp.shape[0]
+                    for i in range(dd):
+                        pos = i / dd * ss
+                        if pos >= ss - 1:
+                            tmp[i] = src[ss - 1]
+                        else:
+                            p0 = int(math.floor(pos))
+                            ii = pos - p0
+                            tmp[i] = src[p0] * (1 - ii) + src[p0 + 1] * (ii)
+                    mm[k] = tmp.reshape(mm[k].shape)
+                    sss = src.squeeze().float().cpu().numpy()
+                    print(sss[:10], "...", sss[-10:])
+                    mmm = mm[k].squeeze().float().cpu().numpy()
+                    print(mmm[:10], "...", mmm[-10:])
+
+    print(f"Save to {init_weight_name}...")
+    torch.save(mm, init_weight_name)
+
+    if model.args.my_pile_stage == 1:
+        print("Done. Now go for stage 2.")
+        exit(0)

finetune/lora/v6/src/utils.py

@@ -0,0 +1,139 @@
+import json, time, random, os
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+time_slot = {}
+time_ref = time.time_ns()
+
+
+def record_time(name):
+    if name not in time_slot:
+        time_slot[name] = 1e20
+    tt = (time.time_ns() - time_ref) / 1e9
+    if tt < time_slot[name]:
+        time_slot[name] = tt
+
+
+class TOKENIZER:
+    def __init__(self, WORD_NAME, UNKNOWN_CHAR="\ue083"):
+        if "list" in str(type(WORD_NAME)):
+            self.charMode = False
+            if WORD_NAME[0] == WORD_NAME[1]:
+                from transformers import PreTrainedTokenizerFast
+
+                self.tokenizer = PreTrainedTokenizerFast(tokenizer_file=WORD_NAME[0])
+            else:
+                from transformers import GPT2TokenizerFast
+
+                self.tokenizer = GPT2TokenizerFast(WORD_NAME[0], WORD_NAME[1])
+            self.vocab_size = len(self.tokenizer)
+        else:
+            self.charMode = True
+            with open(WORD_NAME + ".json", "r", encoding="utf-16") as result_file:
+                self.word_table = json.load(result_file)
+
+            self.vocab_size = len(self.word_table)
+
+            self.stoi = {v: int(k) for k, v in self.word_table.items()}
+            self.itos = {int(k): v for k, v in self.word_table.items()}
+
+            self.UNKNOWN_CHAR = self.stoi[UNKNOWN_CHAR]
+
+    def refine_context(self, context):
+        context = context.strip().split("\n")
+        for c in range(len(context)):
+            context[c] = context[c].strip().strip("\u3000").strip("\r")
+        context = list(filter(lambda c: c != "", context))
+        context = "\n" + ("\n".join(context)).strip()
+        if context == "":
+            context = "\n"
+        return context
+
+    def sample_logits(
+        self, out, x, ctx_len, temperature=1.0, top_p_usual=None, top_p_newline=None
+    ):
+        # out[self.UNKNOWN_CHAR] = -float('Inf')
+        lastChar = int(x[-1])
+
+        probs = F.softmax(out, dim=-1)
+
+        if self.charMode:
+            if self.itos[lastChar] == "\n":
+                top_p = top_p_newline
+            else:
+                top_p = top_p_usual
+        else:
+            top_p = top_p_usual
+
+        if os.environ["RWKV_RUN_DEVICE"] == "cpu":
+            probs = probs.numpy()
+            sorted_probs = np.sort(probs)[::-1]
+            cumulative_probs = np.cumsum(sorted_probs)
+            cutoff = float(sorted_probs[np.argmax(cumulative_probs > top_p)])
+            probs[probs < cutoff] = 0
+            if temperature != 1.0:
+                probs = probs.pow(1.0 / temperature)
+            probs = probs / np.sum(probs)
+            out = np.random.choice(a=len(probs), p=probs)
+            return out
+        else:
+            sorted_probs = torch.sort(probs, descending=True)[0]
+            cumulative_probs = torch.cumsum(sorted_probs, dim=-1).cpu().numpy()
+            cutoff = float(sorted_probs[np.argmax(cumulative_probs > top_p)])
+            probs[probs < cutoff] = 0
+            if temperature != 1.0:
+                probs = probs.pow(1.0 / temperature)
+            out = torch.multinomial(probs, num_samples=1)[0]
+            return out
+
+
+def MaybeIsPrime(number):
+    if FermatPrimalityTest(number) and MillerRabinPrimalityTest(number):
+        return True
+    else:
+        return False
+
+
+def FermatPrimalityTest(number):
+    if number > 1:
+        for time in range(3):
+            randomNumber = random.randint(2, number) - 1
+            if pow(randomNumber, number - 1, number) != 1:
+                return False
+        return True
+    else:
+        return False
+
+
+def MillerRabinPrimalityTest(number):
+    if number == 2:
+        return True
+    elif number == 1 or number % 2 == 0:
+        return False
+    oddPartOfNumber = number - 1
+    timesTwoDividNumber = 0
+    while oddPartOfNumber % 2 == 0:
+        oddPartOfNumber = oddPartOfNumber // 2
+        timesTwoDividNumber = timesTwoDividNumber + 1
+
+    for time in range(3):
+        while True:
+            randomNumber = random.randint(2, number) - 1
+            if randomNumber != 0 and randomNumber != 1:
+                break
+
+        randomNumberWithPower = pow(randomNumber, oddPartOfNumber, number)
+
+        if (randomNumberWithPower != 1) and (randomNumberWithPower != number - 1):
+            iterationNumber = 1
+
+            while (iterationNumber <= timesTwoDividNumber - 1) and (
+                randomNumberWithPower != number - 1
+            ):
+                randomNumberWithPower = pow(randomNumberWithPower, 2, number)
+                iterationNumber = iterationNumber + 1
+            if randomNumberWithPower != (number - 1):
+                return False
+
+    return True

finetune/lora/v6/train.py

@@ -0,0 +1,436 @@
+########################################################################################################
+# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM
+########################################################################################################
+
+import logging
+
+logging.basicConfig(level=logging.INFO)
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+    from pytorch_lightning import Trainer
+    from pytorch_lightning.utilities import rank_zero_info, rank_zero_only
+    import pytorch_lightning as pl
+
+    rank_zero_info("########## work in progress ##########")
+
+    parser = ArgumentParser()
+
+    parser.add_argument("--load_model", default="", type=str)  # full path, with .pth
+    parser.add_argument(
+        "--wandb", default="", type=str
+    )  # wandb project name. if "" then don't use wandb
+    parser.add_argument("--proj_dir", default="out", type=str)
+    parser.add_argument("--random_seed", default="-1", type=int)
+
+    parser.add_argument("--data_file", default="", type=str)
+    parser.add_argument("--data_type", default="utf-8", type=str)
+    parser.add_argument(
+        "--vocab_size", default=0, type=int
+    )  # vocab_size = 0 means auto (for char-level LM and .txt data)
+
+    parser.add_argument("--ctx_len", default=1024, type=int)
+    parser.add_argument(
+        "--epoch_steps", default=1000, type=int
+    )  # a mini "epoch" has [epoch_steps] steps
+    parser.add_argument(
+        "--epoch_count", default=500, type=int
+    )  # train for this many "epochs". will continue afterwards with lr = lr_final
+    parser.add_argument(
+        "--epoch_begin", default=0, type=int
+    )  # if you load a model trained for x "epochs", set epoch_begin = x
+    parser.add_argument(
+        "--epoch_save", default=5, type=int
+    )  # save the model every [epoch_save] "epochs"
+
+    parser.add_argument(
+        "--micro_bsz", default=12, type=int
+    )  # micro batch size (batch size per GPU)
+    parser.add_argument("--n_layer", default=6, type=int)
+    parser.add_argument("--n_embd", default=512, type=int)
+    parser.add_argument("--dim_att", default=0, type=int)
+    parser.add_argument("--dim_ffn", default=0, type=int)
+    parser.add_argument(
+        "--pre_ffn", default=0, type=int
+    )  # replace first att layer by ffn (sometimes better)
+    parser.add_argument("--head_qk", default=0, type=int)  # my headQK trick
+    parser.add_argument("--tiny_att_dim", default=0, type=int)  # tiny attention dim
+    parser.add_argument(
+        "--tiny_att_layer", default=-999, type=int
+    )  # tiny attention @ which layer
+
+    parser.add_argument(
+        "--lr_init", default=6e-4, type=float
+    )  # 6e-4 for L12-D768, 4e-4 for L24-D1024, 3e-4 for L24-D2048
+    parser.add_argument("--lr_final", default=1e-5, type=float)
+    parser.add_argument(
+        "--warmup_steps", default=-1, type=int
+    )  # try 50 if you load a model
+    parser.add_argument("--beta1", default=0.9, type=float)
+    parser.add_argument(
+        "--beta2", default=0.99, type=float
+    )  # use 0.999 when your model is close to convergence
+    parser.add_argument("--adam_eps", default=1e-8, type=float)
+    parser.add_argument(
+        "--grad_cp", default=0, type=int
+    )  # gradient checkpt: saves VRAM, but slower
+    parser.add_argument(
+        "--dropout", default=0, type=float
+    )  # try 0.01 / 0.02 / 0.05 / 0.1
+    parser.add_argument(
+        "--weight_decay", default=0, type=float
+    )  # try 0.1 / 0.01 / 0.001
+    parser.add_argument("--weight_decay_final", default=-1, type=float)
+
+    parser.add_argument(
+        "--my_pile_version", default=1, type=int
+    )  # my special pile version
+    parser.add_argument("--my_pile_stage", default=0, type=int)  # my special pile mode
+    parser.add_argument(
+        "--my_pile_shift", default=-1, type=int
+    )  # my special pile mode - text shift
+    parser.add_argument("--my_pile_edecay", default=0, type=int)
+    parser.add_argument(
+        "--layerwise_lr", default=1, type=int
+    )  # layerwise lr for faster convergence (but slower it/s)
+    parser.add_argument(
+        "--ds_bucket_mb", default=200, type=int
+    )  # deepspeed bucket size in MB. 200 seems enough
+    # parser.add_argument("--cuda_cleanup", default=0, type=int)  # extra cuda cleanup (sometimes helpful)
+
+    parser.add_argument("--my_sample_len", default=0, type=int)
+    parser.add_argument("--my_ffn_shift", default=1, type=int)
+    parser.add_argument("--my_att_shift", default=1, type=int)
+    parser.add_argument(
+        "--head_size_a", default=64, type=int
+    )  # can try larger values for larger models
+    parser.add_argument("--head_size_divisor", default=8, type=int)
+    parser.add_argument("--my_pos_emb", default=0, type=int)
+    parser.add_argument("--load_partial", default=0, type=int)
+    parser.add_argument("--magic_prime", default=0, type=int)
+    parser.add_argument("--my_qa_mask", default=0, type=int)
+    parser.add_argument("--my_random_steps", default=0, type=int)
+    parser.add_argument("--my_testing", default="", type=str)
+    parser.add_argument("--my_exit", default=99999999, type=int)
+    parser.add_argument("--my_exit_tokens", default=0, type=int)
+
+    # LORA
+    parser.add_argument("--emb", action="store_true")
+    parser.add_argument("--lora", action="store_true")
+    parser.add_argument("--lora_load", default="", type=str)
+    parser.add_argument("--lora_r", default=8, type=int)
+    parser.add_argument("--lora_alpha", default=32, type=float)
+    parser.add_argument("--lora_dropout", default=0.01, type=float)
+    parser.add_argument("--lora_parts", default="att,ln,time", type=str)
+
+    if pl.__version__[0] == "2":
+        parser.add_argument("--accelerator", default="gpu", type=str)
+        parser.add_argument("--strategy", default="auto", type=str)
+        parser.add_argument("--devices", default=1, type=int)
+        parser.add_argument("--num_nodes", default=1, type=int)
+        parser.add_argument("--precision", default="fp16", type=str)
+        parser.add_argument("--accumulate_grad_batches", default=1, type=int)
+    else:
+        parser = Trainer.add_argparse_args(parser)
+    args = parser.parse_args()
+
+    ########################################################################################################
+
+    import os, warnings, math, datetime, sys, time
+    import numpy as np
+    import torch
+    from torch.utils.data import DataLoader
+
+    if "deepspeed" in args.strategy:
+        import deepspeed
+    from pytorch_lightning import seed_everything
+
+    if args.random_seed >= 0:
+        print(
+            f"########## WARNING: GLOBAL SEED {args.random_seed} THIS WILL AFFECT MULTIGPU SAMPLING ##########\n"
+            * 3
+        )
+        seed_everything(args.random_seed)
+
+    np.set_printoptions(precision=4, suppress=True, linewidth=200)
+    warnings.filterwarnings(
+        "ignore", ".*Consider increasing the value of the `num_workers` argument*"
+    )
+    warnings.filterwarnings(
+        "ignore", ".*The progress bar already tracks a metric with the*"
+    )
+    # os.environ["WDS_SHOW_SEED"] = "1"
+
+    args.my_timestamp = datetime.datetime.today().strftime("%Y-%m-%d-%H-%M-%S")
+    args.enable_checkpointing = False
+    args.replace_sampler_ddp = False
+    args.logger = False
+    args.gradient_clip_val = 1.0
+    args.num_sanity_val_steps = 0
+    args.check_val_every_n_epoch = int(1e20)
+    args.log_every_n_steps = int(1e20)
+    args.max_epochs = args.epoch_count  # -1 continue forever
+    args.betas = (args.beta1, args.beta2)
+    args.real_bsz = int(args.num_nodes) * int(args.devices) * args.micro_bsz
+    os.environ["RWKV_MY_TESTING"] = args.my_testing
+    os.environ["RWKV_CTXLEN"] = str(args.ctx_len)
+    os.environ["RWKV_HEAD_SIZE_A"] = str(args.head_size_a)
+    if args.dim_att <= 0:
+        args.dim_att = args.n_embd
+    if args.dim_ffn <= 0:
+        args.dim_ffn = int((args.n_embd * 3.5) // 32 * 32)  # default = 3.5x emb size
+
+    if args.data_type == "wds_img":
+        args.run_name = f"v{args.my_img_version}-{args.my_img_size}-{args.my_img_bit}bit-{args.my_img_clip}x{args.my_img_clip_scale}"
+        args.proj_dir = f"{args.proj_dir}-{args.run_name}"
+    else:
+        args.run_name = (
+            f"{args.vocab_size} ctx{args.ctx_len} L{args.n_layer} D{args.n_embd}"
+        )
+    if not os.path.exists(args.proj_dir):
+        os.makedirs(args.proj_dir)
+
+    if args.my_pile_stage > 0:
+        magic_prime_bak = args.magic_prime
+
+        if args.my_pile_shift < 0:
+            args.my_pile_shift = 0
+
+        if magic_prime_bak > 0:
+            args.magic_prime = magic_prime_bak
+        if args.my_qa_mask == 2:
+            args.epoch_count = 2 * args.magic_prime // 40320
+        else:
+            args.epoch_count = args.magic_prime // 40320
+
+        args.epoch_steps = 40320 // args.real_bsz
+        assert args.epoch_steps * args.real_bsz == 40320
+        # if args.my_pile_stage == 2:
+        #     assert args.lr_final == args.lr_init
+        if args.my_pile_stage >= 2:  # find latest saved model
+            list_p = []
+            for p in os.listdir(args.proj_dir):
+                if p.startswith("rwkv") and p.endswith(".pth"):
+                    p = ((p.split("-"))[1].split("."))[0]
+                    if p != "final":
+                        if p == "init":
+                            p = -1
+                        else:
+                            p = int(p)
+                        list_p += [p]
+            list_p.sort()
+            max_p = list_p[-1]
+            if len(list_p) > 1:
+                args.my_pile_prev_p = list_p[-2]  # in case max_p is corrupted
+            if max_p == -1:
+                args.load_model = f"{args.proj_dir}/rwkv-init.pth"
+            else:
+                args.load_model = f"{args.proj_dir}/rwkv-{max_p}.pth"
+                if args.warmup_steps < 0:
+                    if args.my_pile_stage == 2:
+                        args.warmup_steps = 10
+                    else:
+                        args.warmup_steps = 30
+            args.epoch_begin = max_p + 1
+
+    samples_per_epoch = args.epoch_steps * args.real_bsz
+    tokens_per_epoch = samples_per_epoch * args.ctx_len
+    try:
+        deepspeed_version = deepspeed.__version__
+    except:
+        deepspeed_version = None
+        pass
+    rank_zero_info(
+        f"""
+############################################################################
+#
+# RWKV-5 {args.precision.upper()} on {args.num_nodes}x{args.devices} {args.accelerator.upper()}, bsz {args.num_nodes}x{args.devices}x{args.micro_bsz}={args.real_bsz}, {args.strategy} {'with grad_cp' if args.grad_cp > 0 else ''}
+#
+# Data = {args.data_file} ({args.data_type}), ProjDir = {args.proj_dir}
+#
+# Epoch = {args.epoch_begin} to {args.epoch_begin + args.epoch_count - 1}, save every {args.epoch_save} epoch
+#
+# Each "epoch" = {args.epoch_steps} steps, {samples_per_epoch} samples, {tokens_per_epoch} tokens
+#
+# Model = {args.n_layer} n_layer, {args.n_embd} n_embd, {args.ctx_len} ctx_len
+#
+# Adam = lr {args.lr_init} to {args.lr_final}, warmup {args.warmup_steps} steps, beta {args.betas}, eps {args.adam_eps}
+#
+# Found torch {torch.__version__}, recommend 1.13.1+cu117 or newer
+# Found deepspeed {deepspeed_version}, recommend 0.7.0 (faster than newer versions)
+# Found pytorch_lightning {pl.__version__}, recommend 1.9.5
+#
+############################################################################
+"""
+    )
+    rank_zero_info(str(vars(args)) + "\n")
+
+    assert args.data_type in ["utf-8", "utf-16le", "numpy", "binidx", "dummy", "uint16"]
+
+    if args.lr_final == 0 or args.lr_init == 0:
+        rank_zero_info(
+            "\n\nNote: lr_final = 0 or lr_init = 0. Using linear LR schedule instead.\n\n"
+        )
+
+    assert args.precision in ["fp32", "tf32", "fp16", "bf16"]
+    os.environ["RWKV_FLOAT_MODE"] = args.precision
+    if args.precision == "fp32":
+        for i in range(10):
+            rank_zero_info(
+                "\n\nNote: you are using fp32 (very slow). Try bf16 / tf32 for faster training.\n\n"
+            )
+    if args.precision == "fp16":
+        rank_zero_info(
+            "\n\nNote: you are using fp16 (might overflow). Try bf16 / tf32 for stable training.\n\n"
+        )
+
+    os.environ["RWKV_JIT_ON"] = "0"
+    if "deepspeed_stage_3" in args.strategy:
+        os.environ["RWKV_JIT_ON"] = "0"
+
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cudnn.enabled = True
+    if args.precision == "fp32":
+        torch.backends.cudnn.allow_tf32 = False
+        torch.backends.cuda.matmul.allow_tf32 = False
+    else:
+        torch.backends.cudnn.allow_tf32 = True
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+    if "32" in args.precision:
+        args.precision = 32
+    elif args.precision == "fp16":
+        args.precision = 16
+    else:
+        args.precision = "bf16"
+
+    ########################################################################################################
+
+    from src.trainer import train_callback, generate_init_weight
+    from src.dataset import MyDataset
+
+    train_data = MyDataset(args)
+    args.vocab_size = train_data.vocab_size
+
+    from src.model import RWKV, LORA_CONFIG, LoraLinear
+
+    if args.lora:
+        assert args.lora_r > 0, "LoRA should have its `r` > 0"
+        LORA_CONFIG["r"] = args.lora_r
+        LORA_CONFIG["alpha"] = args.lora_alpha
+        LORA_CONFIG["dropout"] = args.lora_dropout
+        LORA_CONFIG["parts"] = set(str(args.lora_parts).split(","))
+        enable_time_finetune = "time" in LORA_CONFIG["parts"]
+        enable_ln_finetune = "ln" in LORA_CONFIG["parts"]
+    model = RWKV(args)
+
+    if args.lora:
+        model.requires_grad_(False)
+        for name, module in model.named_modules():
+
+            if any(n.startswith("lora_") for n, _ in module.named_parameters()):
+                print(f"  LoRA additionally training module {name}")
+                for pname, param in module.named_parameters():
+                    param.requires_grad = "lora_" in pname
+            elif enable_ln_finetune and ".ln" in name:
+                print(f"  LoRA additionally training module {name}")
+                for param in module.parameters():
+                    param.requires_grad = True
+            elif enable_time_finetune and any(
+                n.startswith("time") for n, _ in module.named_parameters()
+            ):
+                for pname, param in module.named_parameters():
+                    if pname.startswith("time"):
+                        print(f"  LoRA additionally training parameter {pname}")
+                        param.requires_grad = True
+
+    if (
+        len(args.load_model) == 0 or args.my_pile_stage == 1
+    ):  # shall we build the initial weights?
+        init_weight_name = f"{args.proj_dir}/rwkv-init.pth"
+        generate_init_weight(model, init_weight_name)  # save initial weights
+        args.load_model = init_weight_name
+
+    rank_zero_info(f"########## Loading {args.load_model}... ##########")
+    try:
+        load_dict = torch.load(args.load_model, map_location="cpu")
+        load_keys = list(load_dict.keys())
+        for k in load_keys:
+            if k.startswith("_forward_module."):
+                load_dict[k.replace("_forward_module.", "")] = load_dict[k]
+                del load_dict[k]
+    except:
+        rank_zero_info(f"Bad checkpoint {args.load_model}")
+        if args.my_pile_stage >= 2:  # try again using another checkpoint
+            max_p = args.my_pile_prev_p
+            if max_p == -1:
+                args.load_model = f"{args.proj_dir}/rwkv-init.pth"
+            else:
+                args.load_model = f"{args.proj_dir}/rwkv-{max_p}.pth"
+            args.epoch_begin = max_p + 1
+            rank_zero_info(f"Trying {args.load_model}")
+            load_dict = torch.load(args.load_model, map_location="cpu")
+
+    if args.load_partial == 1:
+        load_keys = load_dict.keys()
+        for k in model.state_dict():
+            if k not in load_keys:
+                load_dict[k] = model.state_dict()[k]
+    model.load_state_dict(load_dict, strict=(not args.lora))
+    if os.path.isfile(args.lora_load):
+        model.load_state_dict(
+            torch.load(args.lora_load, map_location="cpu"), strict=False
+        )
+
+    if pl.__version__[0] == "2":
+        trainer = Trainer(
+            accelerator=args.accelerator,
+            strategy=args.strategy,
+            devices=args.devices,
+            num_nodes=args.num_nodes,
+            precision=args.precision,
+            logger=args.logger,
+            callbacks=[train_callback(args)],
+            max_epochs=args.max_epochs,
+            check_val_every_n_epoch=args.check_val_every_n_epoch,
+            num_sanity_val_steps=args.num_sanity_val_steps,
+            log_every_n_steps=args.log_every_n_steps,
+            enable_checkpointing=args.enable_checkpointing,
+            accumulate_grad_batches=args.accumulate_grad_batches,
+            gradient_clip_val=args.gradient_clip_val,
+        )
+    else:
+        trainer = Trainer.from_argparse_args(
+            args,
+            callbacks=[train_callback(args)],
+        )
+
+    if trainer.global_rank == 0:
+        for n in model.state_dict():
+            shape = model.state_dict()[n].shape
+            shape = [i for i in shape if i != 1]
+            if len(shape) > 1:
+                print(f"{str(shape[0]).ljust(5)} {str(shape[1]).ljust(5)} {n}")
+            else:
+                print(f"{str(shape[0]).ljust(5)}       {n}")
+
+    if "deepspeed" in args.strategy:
+        trainer.strategy.config["zero_optimization"]["allgather_bucket_size"] = (
+            args.ds_bucket_mb * 1000 * 1000
+        )
+        trainer.strategy.config["zero_optimization"]["reduce_bucket_size"] = (
+            args.ds_bucket_mb * 1000 * 1000
+        )
+
+    # must set shuffle=False, persistent_workers=False (because worker is in another thread)
+    data_loader = DataLoader(
+        train_data,
+        shuffle=False,
+        pin_memory=True,
+        batch_size=args.micro_bsz,
+        num_workers=1,
+        persistent_workers=False,
+        drop_last=True,
+    )
+
+    trainer.fit(model, data_loader)

frontend/package.json

@@ -9,7 +9,7 @@
     "preview": "vite preview"
   },
   "dependencies": {
-    "@fluentui/react-components": "^9.20.0",
+    "@fluentui/react-components": "^9.47.2",
     "@fluentui/react-icons": "^2.0.201",
     "@magenta/music": "^1.23.1",
     "@microsoft/fetch-event-source": "^2.0.1",
@@ -20,6 +20,7 @@
     "file-saver": "^2.0.5",
     "html-midi-player": "^1.5.0",
     "i18next": "^22.4.15",
+    "katex": "^0.16.9",
     "lodash-es": "^4.17.21",
     "mobx": "^6.9.0",
     "mobx-react-lite": "^3.4.3",
@@ -35,13 +36,16 @@
     "react-router-dom": "^6.11.1",
     "react-toastify": "^9.1.3",
     "rehype-highlight": "^6.0.0",
+    "rehype-katex": "^6.0.3",
     "rehype-raw": "^6.1.1",
     "remark-breaks": "^3.0.3",
     "remark-gfm": "^3.0.1",
+    "remark-math": "^5.1.1",
     "usehooks-ts": "^2.9.1",
     "uuid": "^9.0.0"
   },
   "devDependencies": {
+    "@tailwindcss/typography": "^0.5.10",
     "@types/file-saver": "^2.0.7",
     "@types/lodash-es": "^4.17.12",
     "@types/react": "^18.2.6",

frontend/src/App.tsx

@@ -83,7 +83,7 @@ const App: FC = observer(() => {
             ))}
           </TabList>
         </div>
-        <div className="h-full w-full p-2 box-border overflow-y-hidden">
+        <div className="h-full w-full py-2 pr-2 box-border overflow-y-hidden">
           <Routes>
             {pages.map(({ path, element }, index) => (
               <Route key={`${path}-${index}`} path={path} element={<LazyImportComponent lazyChildren={element} />} />

frontend/src/_locales/ja/main.json

@@ -4,7 +4,7 @@
   "About": "約",
   "Settings": "設定",
   "Go to chat page": "チャットページに移動する",
-  "Manage your configs": "あなたの設定を管理する",
+  "Manage your configs, adjust the starting model and parameters": "設定を管理し、開始モデルとパラメータを調整します",
   "Manage models": "モデルの管理",
   "Run": "実行",
   "Offline": "オフライン",
@@ -96,7 +96,7 @@
   "Python dependencies are incomplete, would you like to install them?": "Pythonの依存関係が不完全です、インストールしますか？",
   "Install": "インストール",
   "This is the latest version": "これは最新バージョンです",
-  "Use Tsinghua Pip Mirrors": "清華大学Pipミラーサーバーを使用",
+  "Use Alibaba Cloud Pip Mirrors": "Alibaba Cloud Pipミラーサーバーを使用",
   "Model Config Exception": "モデル設定例外",
   "Use Gitee Updates Source": "Gitee更新ソースを使用",
   "Use Custom CUDA kernel to Accelerate": "カスタムCUDAカーネルを使用して加速",
@@ -347,5 +347,12 @@
   "Parallel Token Chunk Size": "並列トークンチャンクサイズ",
   "Maximum tokens to be processed in parallel at once. For high end GPUs, this could be 64 or 128 (faster).": "一度に並列で処理される最大トークン数。高性能なGPUの場合、64または128になります（高速）。",
   "Global Penalty": "グローバルペナルティ",
-  "When generating a response, whether to include the submitted prompt as a penalty factor. By turning this off, you will get the same generated results as official RWKV Gradio. If you find duplicate results in the generated results, turning this on can help avoid generating duplicates.": "レスポンスを生成する際、提出されたプロンプトをペナルティ要因として含めるかどうか。これをオフにすると、公式RWKV Gradioと同じ生成結果を得ることができます。生成された結果に重複がある場合、これをオンにすることで重複の生成を回避するのに役立ちます。"
+  "When generating a response, whether to include the submitted prompt as a penalty factor. By turning this off, you will get the same generated results as official RWKV Gradio. If you find duplicate results in the generated results, turning this on can help avoid generating duplicates.": "レスポンスを生成する際、提出されたプロンプトをペナルティ要因として含めるかどうか。これをオフにすると、公式RWKV Gradioと同じ生成結果を得ることができます。生成された結果に重複がある場合、これをオンにすることで重複の生成を回避するのに役立ちます。",
+  "Create a new user or AI message content. You can prepare a chat record with AI here, and fill in the responses you want to get from AI in the tone of AI. When you use this preset, the chat record will be processed, and at this point, AI will better understand what you want it to do or what role to play.": "新しいユーザーまたはAIメッセージコンテンツを作成します。ここでAIとのチャット記録を準備し、AIから得たい応答をAIのトーンで記入することができます。このプリセットを使用すると、チャット記録が処理され、この時点でAIはあなたが望むことやどのような役割を果たすかをよりよく理解することができます。",
+  "The name used internally by the model when processing user message, changing this value helps improve the role-playing effect.": "ユーザーメッセージを処理する際にモデルが内部で使用する名前、この値を変更することで、役割演技の効果を向上させることができます。",
+  "The name used internally by the model when processing AI message, changing this value helps improve the role-playing effect.": "AIメッセージを処理する際にモデルが内部で使用する名前、この値を変更することで、役割演技の効果を向上させることができます。",
+  "Inside the model, there is a default prompt to improve the model's handling of common issues, but it may degrade the role-playing effect. You can disable this option to achieve a better role-playing effect.": "モデル内部には、一般的な問題の処理を改善するためのデフォルトのプロンプトがありますが、役割演技の効果を低下させる可能性があります。このオプションを無効にすることで、より良い役割演技効果を得ることができます。",
+  "Exit without saving": "保存せずに終了",
+  "Content has been changed, are you sure you want to exit without saving?": "コンテンツが変更されています、保存せずに終了してもよろしいですか？",
+  "Don't forget to correctly fill in your Ollama API Chat Model Name.": "Ollama APIチャットモデル名を正しく記入するのを忘れないでください。"
 }

frontend/src/_locales/zh-hans/main.json

@@ -4,7 +4,7 @@
   "About": "关于",
   "Settings": "设置",
   "Go to chat page": "前往聊天页",
-  "Manage your configs": "管理你的配置",
+  "Manage your configs, adjust the starting model and parameters": "管理你的配置, 调整启动的模型和参数",
   "Manage models": "管理模型",
   "Run": "运行",
   "Offline": "离线",
@@ -96,7 +96,7 @@
   "Python dependencies are incomplete, would you like to install them?": "Python依赖缺失, 是否安装?",
   "Install": "安装",
   "This is the latest version": "已是最新版",
-  "Use Tsinghua Pip Mirrors": "使用清华大学Pip镜像源",
+  "Use Alibaba Cloud Pip Mirrors": "使用阿里云Pip镜像源",
   "Model Config Exception": "模型配置异常",
   "Use Gitee Updates Source": "使用Gitee更新源",
   "Use Custom CUDA kernel to Accelerate": "使用自定义CUDA算子加速",
@@ -347,5 +347,12 @@
   "Parallel Token Chunk Size": "并行Token块大小",
   "Maximum tokens to be processed in parallel at once. For high end GPUs, this could be 64 or 128 (faster).": "一次最多可以并行处理的token数量. 对于高端显卡, 这可以是64或128 (更快)",
   "Global Penalty": "全局惩罚",
-  "When generating a response, whether to include the submitted prompt as a penalty factor. By turning this off, you will get the same generated results as official RWKV Gradio. If you find duplicate results in the generated results, turning this on can help avoid generating duplicates.": "生成响应时, 是否将提交的prompt也纳入到惩罚项. 关闭此项将得到与RWKV官方Gradio完全一致的生成结果. 如果你发现生成结果出现重复, 那么开启此项有助于避免生成重复"
+  "When generating a response, whether to include the submitted prompt as a penalty factor. By turning this off, you will get the same generated results as official RWKV Gradio. If you find duplicate results in the generated results, turning this on can help avoid generating duplicates.": "生成响应时, 是否将提交的prompt也纳入到惩罚项. 关闭此项将得到与RWKV官方Gradio完全一致的生成结果. 如果你发现生成结果出现重复, 那么开启此项有助于避免生成重复",
+  "Create a new user or AI message content. You can prepare a chat record with AI here, and fill in the responses you want to get from AI in the tone of AI. When you use this preset, the chat record will be processed, and at this point, AI will better understand what you want it to do or what role to play.": "新建一个 用户 或 AI 的发言内容. 你可以在这里准备好一段你与 AI 的聊天记录, 并用 AI 的口吻正确填写你想得到的 AI 的回复, 这样你在使用这个预设时, 这些聊天记录也会被处理, 并且此时 AI 能更好地理解你希望它做什么 / 扮演什么样的角色",
+  "The name used internally by the model when processing user message, changing this value helps improve the role-playing effect.": "模型内部处理用户发言时使用的名称, 更改此值有助于改善角色扮演效果",
+  "The name used internally by the model when processing AI message, changing this value helps improve the role-playing effect.": "模型内部处理AI发言时使用的名称, 更改此值有助于改善角色扮演效果",
+  "Inside the model, there is a default prompt to improve the model's handling of common issues, but it may degrade the role-playing effect. You can disable this option to achieve a better role-playing effect.": "模型内部有一个默认提示来改善模型处理常规问题的效果, 但它可能会让角色扮演的效果变差, 你可以关闭此选项来获得更好的角色扮演效果",
+  "Exit without saving": "退出而不保存",
+  "Content has been changed, are you sure you want to exit without saving?": "内容已经被修改, 你确定要退出而不保存吗?",
+  "Don't forget to correctly fill in your Ollama API Chat Model Name.": "不要忘记正确填写你的Ollama API 聊天模型名"
 }

frontend/src/components/DialogButton.tsx

@@ -16,20 +16,28 @@ import { LazyImportComponent } from './LazyImportComponent';
 const MarkdownRender = React.lazy(() => import('./MarkdownRender'));
 
 export const DialogButton: FC<{
-  text?: string | null
+  text?: string | null,
   icon?: ReactElement,
   tooltip?: string | null,
   className?: string,
   title: string,
-  contentText: string,
+  content?: string | ReactElement | null
   markdown?: boolean,
   onConfirm?: () => void,
   size?: 'small' | 'medium' | 'large',
   shape?: 'rounded' | 'circular' | 'square',
   appearance?: 'secondary' | 'primary' | 'outline' | 'subtle' | 'transparent',
+  cancelButton?: boolean,
+  confirmButton?: boolean,
+  cancelButtonText?: string,
+  confirmButtonText?: string,
 }> = ({
-  text, icon, tooltip, className, title, contentText, markdown,
+  text, icon, tooltip, className, title, content, markdown,
-  onConfirm, size, shape, appearance
+  onConfirm, size, shape, appearance,
+  cancelButton = true,
+  confirmButton = true,
+  cancelButtonText = 'Cancel',
+  confirmButtonText = 'Confirm'
 }) => {
   const { t } = useTranslation();
 
@@ -41,26 +49,31 @@ export const DialogButton: FC<{
         <Button className={className} icon={icon} size={size} shape={shape} appearance={appearance}>{text}</Button>
       }
     </DialogTrigger>
-    <DialogSurface>
+    <DialogSurface style={{ transform: 'unset' }}>
       <DialogBody>
         <DialogTitle>{title}</DialogTitle>
         <DialogContent>
           {
             markdown ?
               <LazyImportComponent lazyChildren={MarkdownRender}>
-                {contentText}
+                {content}
               </LazyImportComponent> :
-              contentText
+              content
           }
         </DialogContent>
         <DialogActions>
-          <DialogTrigger disableButtonEnhancement>
+          {cancelButton && (
-            <Button appearance="secondary">{t('Cancel')}</Button>
+            <DialogTrigger disableButtonEnhancement>
-          </DialogTrigger>
+              <Button appearance="secondary">{t(cancelButtonText)}</Button>
-          <DialogTrigger disableButtonEnhancement>
+            </DialogTrigger>
-            <Button appearance="primary" onClick={onConfirm}>{t('Confirm')}
+          )}
-            </Button>
+          {confirmButton && (
-          </DialogTrigger>
+            <DialogTrigger disableButtonEnhancement>
+              <Button appearance="primary" onClick={onConfirm}>
+                {t(confirmButtonText)}
+              </Button>
+            </DialogTrigger>
+          )}
         </DialogActions>
       </DialogBody>
     </DialogSurface>

frontend/src/components/MarkdownRender.tsx

@@ -1,6 +1,9 @@
+import 'katex/dist/katex.min.css';
 import ReactMarkdown from 'react-markdown';
 import rehypeRaw from 'rehype-raw';
 import rehypeHighlight from 'rehype-highlight';
+import rehypeKatex from 'rehype-katex';
+import remarkMath from 'remark-math';
 import remarkGfm from 'remark-gfm';
 import remarkBreaks from 'remark-breaks';
 import { FC } from 'react';
@@ -23,12 +26,12 @@ const Hyperlink: FC<any> = ({ href, children }) => {
 
 const MarkdownRender: FC<ReactMarkdownOptions & { disabled?: boolean }> = (props) => {
   return (
-    <div dir="auto" className="markdown-body">
+    <div dir="auto" className="prose markdown-body" style={{ maxWidth: '100%' }}>
       {props.disabled ?
-        <div style={{ whiteSpace: 'pre-wrap' }}>
+        <div style={{ whiteSpace: 'pre-wrap' }} className={props.className}>
           {props.children}
         </div> :
-        <ReactMarkdown
+        <ReactMarkdown className={props.className}
           allowedElements={[
             'div',
             'p',
@@ -90,8 +93,9 @@ const MarkdownRender: FC<ReactMarkdownOptions & { disabled?: boolean }> = (props
             'cite'
           ]}
           unwrapDisallowed={true}
-          remarkPlugins={[remarkGfm, remarkBreaks]}
+          remarkPlugins={[remarkMath, remarkGfm, remarkBreaks]}
           rehypePlugins={[
+            rehypeKatex,
             rehypeRaw,
             [
               rehypeHighlight,

frontend/src/components/ResetConfigsButton.tsx

@@ -9,7 +9,7 @@ import { defaultModelConfigs, defaultModelConfigsMac } from '../pages/defaultCon
 export const ResetConfigsButton: FC<{ afterConfirm?: () => void }> = ({ afterConfirm }) => {
   const { t } = useTranslation();
   return <DialogButton icon={<ArrowReset20Regular />} tooltip={t('Reset All Configs')} title={t('Reset All Configs')}
-    contentText={t('Are you sure you want to reset all configs? This will obtain the latest preset configs, but will override your custom configs and cannot be undone.')}
+    content={t('Are you sure you want to reset all configs? This will obtain the latest preset configs, but will override your custom configs and cannot be undone.')}
     onConfirm={() => {
       commonStore.setModelConfigs(commonStore.platform !== 'darwin' ? defaultModelConfigs : defaultModelConfigsMac, false);
       commonStore.setCurrentConfigIndex(0, true);

frontend/src/components/ToolTipButton.tsx

@@ -26,10 +26,12 @@ export const ToolTipButton: FC<{
   onClick,
   showDelay = 0
 }) => {
-  return (
+  return (desc ?
-    <Tooltip content={desc} showDelay={showDelay} hideDelay={0} relationship="label">
+      <Tooltip content={desc} showDelay={showDelay} hideDelay={0} relationship="label">
+        <Button style={style} className={className} disabled={disabled} icon={icon} onClick={onClick} size={size}
+          shape={shape} appearance={appearance}>{text}</Button>
+      </Tooltip> :
       <Button style={style} className={className} disabled={disabled} icon={icon} onClick={onClick} size={size}
         shape={shape} appearance={appearance}>{text}</Button>
-    </Tooltip>
   );
 };

frontend/src/pages/AudiotrackManager/AudiotrackButton.tsx

@@ -28,7 +28,7 @@ export const AudiotrackButton: FC<{
         {t('Open MIDI Input Audio Tracks')}
       </Button>
     </DialogTrigger>
-    <DialogSurface style={{ paddingTop: 0, maxWidth: '90vw', width: 'fit-content' }}>
+    <DialogSurface style={{ paddingTop: 0, maxWidth: '90vw', width: 'fit-content', transform: 'unset' }}>
       <DialogBody>
         <DialogContent className="overflow-hidden">
           <CustomToastContainer />

frontend/src/pages/AudiotrackManager/AudiotrackEditor.tsx

@@ -32,7 +32,7 @@ import {
   flushMidiRecordingContent,
   getMidiRawContentMainInstrument,
   getMidiRawContentTime,
-  getServerRoot,
+  getReqUrl,
   OpenFileDialog,
   refreshTracksTotalTime
 } from '../../utils';
@@ -474,8 +474,13 @@ const AudiotrackEditor: FC<{ setPrompt: (prompt: string) => void }> = observer((
                 OpenFileDialog('*.mid').then(async blob => {
                   const bodyForm = new FormData();
                   bodyForm.append('file_data', blob);
-                  fetch(getServerRoot(commonStore.getCurrentModelConfig().apiParameters.apiPort) + '/midi-to-text', {
+                  const {
+                    url,
+                    headers
+                  } = await getReqUrl(commonStore.getCurrentModelConfig().apiParameters.apiPort, '/midi-to-text');
+                  fetch(url, {
                     method: 'POST',
+                    headers,
                     body: bodyForm
                   }).then(async r => {
                       if (r.status === 200) {
@@ -494,7 +499,7 @@ const AudiotrackEditor: FC<{ setPrompt: (prompt: string) => void }> = observer((
                         commonStore.setTracks(tracks);
                         refreshTracksTotalTime();
                       } else {
-                        toast(r.statusText + '\n' + (await r.text()), {
+                        toast('Failed to fetch - ' + r.status + ' - ' + r.statusText + ' - ' + (await r.text()), {
                           type: 'error'
                         });
                       }

frontend/src/pages/Chat.tsx

@@ -41,19 +41,20 @@ import { OpenFileFolder, OpenOpenFileDialog, OpenSaveFileDialog } from '../../wa
 import {
   absPathAsset,
   bytesToReadable,
-  getServerRoot,
+  getReqUrl,
   newChatConversation,
   OpenFileDialog,
   setActivePreset,
   toastWithButton
 } from '../utils';
 import { useMediaQuery } from 'usehooks-ts';
-import { botName, ConversationMessage, MessageType, Role, userName, welcomeUuid } from '../types/chat';
+import { botName, ConversationMessage, MessageType, Role, systemName, userName, welcomeUuid } from '../types/chat';
 import { Labeled } from '../components/Labeled';
 import { ValuedSlider } from '../components/ValuedSlider';
 import { PresetsButton } from './PresetsManager/PresetsButton';
 import { webOpenOpenFileDialog } from '../utils/web-file-operations';
 import { defaultPenaltyDecay } from './defaultConfigs';
+import { AvatarProps } from '@fluentui/react-avatar';
 
 let chatSseControllers: {
   [id: string]: AbortController
@@ -92,6 +93,18 @@ const MoreUtilsButton: FC<{
   </Menu>;
 });
 
+const HiddenAvatar: FC<AvatarProps> = ({ children, hidden, ...props }) => {
+  if (hidden) {
+    return <div>{children}</div>;
+  } else {
+    return (
+      <Avatar {...props}>
+        {children}
+      </Avatar>
+    );
+  }
+};
+
 const ChatMessageItem: FC<{
   uuid: string,
   onSubmit: (message: string | null, answerId: string | null,
@@ -129,7 +142,9 @@ const ChatMessageItem: FC<{
   return <div
     className={classnames(
       'flex gap-2 mb-2 overflow-hidden',
-      messageItem.side === 'left' ? 'flex-row' : 'flex-row-reverse'
+      messageItem.side === 'left' && 'flex-row',
+      messageItem.side === 'right' && 'flex-row-reverse',
+      messageItem.side === 'center' && 'flex-row justify-center'
     )}
     onMouseEnter={() => {
       const utils = document.getElementById('utils-' + uuid);
@@ -140,22 +155,26 @@ const ChatMessageItem: FC<{
       if (utils) utils.classList.add('invisible');
     }}
   >
-    <Avatar
+    <HiddenAvatar
       color={messageItem.color}
       name={messageItem.sender}
       image={avatarImg ? { src: avatarImg } : undefined}
+      hidden={messageItem.side === 'center'}
     />
     <div
       className={classnames(
-        'flex p-2 rounded-lg overflow-hidden',
+        'flex rounded-lg overflow-hidden',
         editing ? 'grow' : '',
-        messageItem.side === 'left' ? 'bg-gray-200' : 'bg-blue-500',
+        messageItem.side === 'center' ? 'p-1' : 'p-2',
-        messageItem.side === 'left' ? 'text-gray-600' : 'text-white'
+        commonStore.settings.darkMode ? 'bg-neutral-800 border-neutral-600 border-[1px]' : (messageItem.side === 'right' ? 'bg-blue-500' : 'bg-gray-200'),
+        commonStore.settings.darkMode ? 'text-white' : (messageItem.side === 'right' ? 'text-white' : 'text-gray-600'),
+        messageItem.side === 'center' && 'text-opacity-60 bg-opacity-30 border-opacity-30'
       )}
     >
       {!editing ?
         <div className="flex flex-col">
-          <MarkdownRender disabled={!commonStore.chatParams.markdown}>{messageItem.content}</MarkdownRender>
+          <MarkdownRender className={classnames(messageItem.side === 'center' && 'text-xs')}
+            disabled={!commonStore.chatParams.markdown || messageItem.side === 'center'}>{messageItem.content}</MarkdownRender>
           {uuid in commonStore.attachments &&
             <div className="flex grow">
               <div className="grow" />
@@ -375,7 +394,11 @@ const SidePanel: FC = observer(() => {
             return;
           const messageItem = commonStore.conversation[uuid];
           if (messageItem.type !== MessageType.Error) {
-            savedContent += `${messageItem.sender === userName ? user : bot}: ${messageItem.content}\n\n`;
+            if (messageItem.sender === userName) {
+              savedContent += `${user}: ${messageItem.content}\n\n`;
+            } else if (messageItem.sender === botName) {
+              savedContent += `${bot}: ${messageItem.content}\n\n`;
+            }
           }
         });
 
@@ -418,7 +441,7 @@ const ChatPanel: FC = observer(() => {
         [welcomeUuid]: {
           sender: botName,
           type: MessageType.Normal,
-          color: 'colorful',
+          color: 'neutral',
           avatarImg: logo,
           time: new Date().toISOString(),
           content: commonStore.platform === 'web' ? t('Hello, what can I do for you?') : t('Hello! I\'m RWKV, an open-source and commercially usable large language model.'),
@@ -452,7 +475,7 @@ const ChatPanel: FC = observer(() => {
   // if answerId is not null, override the answer with new response;
   // if startUuid is null, start generating api body messages from first message;
   // if endUuid is null, generate api body messages until last message;
-  const onSubmit = useCallback((message: string | null = null, answerId: string | null = null,
+  const onSubmit = useCallback(async (message: string | null = null, answerId: string | null = null,
     startUuid: string | null = null, endUuid: string | null = null, includeEndUuid: boolean = false) => {
     if (message) {
       const newId = uuid();
@@ -499,6 +522,8 @@ const ChatPanel: FC = observer(() => {
         messages.push({ role: 'user', content: messageItem.content });
       } else if (messageItem.done && messageItem.type === MessageType.Normal && messageItem.sender === botName) {
         messages.push({ role: 'assistant', content: messageItem.content });
+      } else if (messageItem.done && messageItem.type === MessageType.Normal && messageItem.sender === systemName) {
+        messages.push({ role: 'system', content: messageItem.content });
       }
     });
 
@@ -509,7 +534,7 @@ const ChatPanel: FC = observer(() => {
     commonStore.conversation[answerId] = {
       sender: botName,
       type: MessageType.Normal,
-      color: 'colorful',
+      color: 'neutral',
       avatarImg: logo,
       time: new Date().toISOString(),
       content: '',
@@ -520,15 +545,27 @@ const ChatPanel: FC = observer(() => {
     commonStore.setConversationOrder(commonStore.conversationOrder);
     setTimeout(scrollToBottom);
     let answer = '';
+    let finished = false;
+    const finish = () => {
+      finished = true;
+      if (answerId! in chatSseControllers)
+        delete chatSseControllers[answerId!];
+      commonStore.conversation[answerId!].done = true;
+      commonStore.conversation[answerId!].content = commonStore.conversation[answerId!].content.trim();
+      commonStore.setConversation(commonStore.conversation);
+      commonStore.setConversationOrder([...commonStore.conversationOrder]);
+    };
     const chatSseController = new AbortController();
     chatSseControllers[answerId] = chatSseController;
+    const { url, headers } = await getReqUrl(port, '/v1/chat/completions', true);
     fetchEventSource( // https://api.openai.com/v1/chat/completions || http://127.0.0.1:${port}/v1/chat/completions
-      getServerRoot(port, true) + '/v1/chat/completions',
+      url,
       {
         method: 'POST',
         headers: {
           'Content-Type': 'application/json',
-          Authorization: `Bearer ${commonStore.settings.apiKey}`
+          Authorization: `Bearer ${commonStore.settings.apiKey}`,
+          ...headers
         },
         body: JSON.stringify({
           messages: messages.slice(-commonStore.chatParams.historyN),
@@ -547,13 +584,8 @@ const ChatPanel: FC = observer(() => {
         signal: chatSseController?.signal,
         onmessage(e) {
           scrollToBottom();
-          if (e.data.trim() === '[DONE]') {
+          if (!finished && e.data.trim() === '[DONE]') {
-            if (answerId! in chatSseControllers)
+            finish();
-              delete chatSseControllers[answerId!];
-            commonStore.conversation[answerId!].done = true;
-            commonStore.conversation[answerId!].content = commonStore.conversation[answerId!].content.trim();
-            commonStore.setConversation(commonStore.conversation);
-            commonStore.setConversationOrder([...commonStore.conversationOrder]);
             return;
           }
           let data;
@@ -566,6 +598,10 @@ const ChatPanel: FC = observer(() => {
           if (data.model)
             commonStore.setLastModelName(data.model);
           if (data.choices && Array.isArray(data.choices) && data.choices.length > 0) {
+            if (!finished && data.choices[0]?.finish_reason) {
+              finish();
+              return;
+            }
             answer += data.choices[0]?.delta?.content || '';
             commonStore.conversation[answerId!].content = answer;
             commonStore.setConversation(commonStore.conversation);
@@ -574,7 +610,12 @@ const ChatPanel: FC = observer(() => {
         },
         async onopen(response) {
           if (response.status !== 200) {
-            commonStore.conversation[answerId!].content += '\n[ERROR]\n```\n' + response.statusText + '\n' + (await response.text()) + '\n```';
+            let errText = await response.text();
+            try {
+              errText = JSON.stringify(JSON.parse(errText), null, 2);
+            } catch (e) {
+            }
+            commonStore.conversation[answerId!].content += '\n[ERROR]\n```\n' + response.status + ' - ' + response.statusText + '\n' + errText + '\n```';
             commonStore.setConversation(commonStore.conversation);
             commonStore.setConversationOrder([...commonStore.conversationOrder]);
             setTimeout(scrollToBottom);
@@ -617,7 +658,7 @@ const ChatPanel: FC = observer(() => {
           <DialogButton tooltip={t('Clear')}
             icon={<Delete28Regular />}
             size={mq ? 'large' : 'small'} shape="circular" appearance="subtle" title={t('Clear')}
-            contentText={t('Are you sure you want to clear the conversation? It cannot be undone.')}
+            content={t('Are you sure you want to clear the conversation? It cannot be undone.')}
             onConfirm={() => {
               if (generating) {
                 for (const id in chatSseControllers) {
@@ -625,7 +666,7 @@ const ChatPanel: FC = observer(() => {
                 }
                 chatSseControllers = {};
               }
-              setActivePreset(commonStore.activePreset);
+              setActivePreset(commonStore.activePreset, commonStore.activePresetIndex);
             }} />
           <div className="relative flex grow">
             <Textarea
@@ -688,8 +729,10 @@ const ChatPanel: FC = observer(() => {
                         const urlPath = `/file-to-text?file_name=${attachmentName}`;
                         const bodyForm = new FormData();
                         bodyForm.append('file_data', blob, attachmentName);
-                        fetch(getServerRoot(port) + urlPath, {
+                        const { url, headers } = await getReqUrl(port, urlPath);
+                        fetch(url, {
                           method: 'POST',
+                          headers,
                           body: bodyForm
                         }).then(async r => {
                             if (r.status === 200) {
@@ -712,7 +755,7 @@ const ChatPanel: FC = observer(() => {
                                   autoClose: 1000
                                 });
                             } else {
-                              toast(r.statusText + '\n' + (await r.text()), {
+                              toast('Failed to fetch - ' + r.status + ' - ' + r.statusText + ' - ' + (await r.text()), {
                                 type: 'error'
                               });
                             }

frontend/src/pages/Completion.tsx

@@ -14,7 +14,7 @@ import { ToolTipButton } from '../components/ToolTipButton';
 import { ArrowSync20Regular } from '@fluentui/react-icons';
 import { defaultPresets } from './defaultConfigs';
 import { CompletionParams, CompletionPreset } from '../types/completion';
-import { getServerRoot } from '../utils';
+import { getReqUrl } from '../utils';
 
 let completionSseController: AbortController | null = null;
 
@@ -68,7 +68,7 @@ const CompletionPanel: FC = observer(() => {
     });
   };
 
-  const onSubmit = (prompt: string) => {
+  const onSubmit = async (prompt: string) => {
     commonStore.setCompletionSubmittedPrompt(prompt);
 
     if (commonStore.status.status === ModelStatus.Offline && !commonStore.settings.apiUrl && commonStore.platform !== 'web') {
@@ -80,14 +80,21 @@ const CompletionPanel: FC = observer(() => {
     prompt += params.injectStart.replaceAll('\\n', '\n');
 
     let answer = '';
+    let finished = false;
+    const finish = () => {
+      finished = true;
+      commonStore.setCompletionGenerating(false);
+    };
     completionSseController = new AbortController();
+    const { url, headers } = await getReqUrl(port, '/v1/completions', true);
     fetchEventSource( // https://api.openai.com/v1/completions || http://127.0.0.1:${port}/v1/completions
-      getServerRoot(port, true) + '/v1/completions',
+      url,
       {
         method: 'POST',
         headers: {
           'Content-Type': 'application/json',
-          Authorization: `Bearer ${commonStore.settings.apiKey}`
+          Authorization: `Bearer ${commonStore.settings.apiKey}`,
+          ...headers
         },
         body: JSON.stringify({
           prompt,
@@ -103,8 +110,8 @@ const CompletionPanel: FC = observer(() => {
         signal: completionSseController?.signal,
         onmessage(e) {
           scrollToBottom();
-          if (e.data.trim() === '[DONE]') {
+          if (!finished && e.data.trim() === '[DONE]') {
-            commonStore.setCompletionGenerating(false);
+            finish();
             return;
           }
           let data;
@@ -117,13 +124,17 @@ const CompletionPanel: FC = observer(() => {
           if (data.model)
             commonStore.setLastModelName(data.model);
           if (data.choices && Array.isArray(data.choices) && data.choices.length > 0) {
+            if (!finished && data.choices[0]?.finish_reason) {
+              finish();
+              return;
+            }
             answer += data.choices[0]?.text || data.choices[0]?.delta?.content || '';
             setPrompt(prompt + answer.replace(/\s+$/, '') + params.injectEnd.replaceAll('\\n', '\n'));
           }
         },
         async onopen(response) {
           if (response.status !== 200) {
-            toast(response.statusText + '\n' + (await response.text()), {
+            toast(response.status + ' - ' + response.statusText + ' - ' + (await response.text()), {
               type: 'error'
             });
           }
@@ -275,7 +286,7 @@ const CompletionPanel: FC = observer(() => {
             onSubmit(commonStore.completionSubmittedPrompt);
           }} />
           <DialogButton className="grow" text={t('Reset')} title={t('Reset')}
-            contentText={t('Are you sure you want to reset this page? It cannot be undone.')}
+            content={t('Are you sure you want to reset this page? It cannot be undone.')}
             onConfirm={() => {
               setPreset(defaultPresets.find((preset) => preset.name === name)!);
             }} />

frontend/src/pages/Composition.tsx

@@ -26,7 +26,7 @@ import {
   SaveFile,
   StartFile
 } from '../../wailsjs/go/backend_golang/App';
-import { getServerRoot, getSoundFont, toastWithButton } from '../utils';
+import { getReqUrl, getSoundFont, toastWithButton } from '../utils';
 import { CompositionParams } from '../types/composition';
 import { useMediaQuery } from 'usehooks-ts';
 import { AudiotrackButton } from './AudiotrackManager/AudiotrackButton';
@@ -135,7 +135,7 @@ const CompositionPanel: FC = observer(() => {
     }
   }, [commonStore.midiPorts]);
 
-  const generateNs = (autoPlay: boolean) => {
+  const generateNs = async (autoPlay: boolean) => {
     if (commonStore.getCurrentModelConfig().modelParameters.modelName.toLowerCase().includes('abc')) {
       import('abcjs').then(ABCJS => {
         ABCJS.renderAbc('abc-paper', commonStore.compositionParams.prompt, { responsive: 'resize' });
@@ -143,10 +143,12 @@ const CompositionPanel: FC = observer(() => {
       return;
     }
 
-    fetch(getServerRoot(port) + '/text-to-midi', {
+    const { url, headers } = await getReqUrl(port, '/text-to-midi');
+    fetch(url, {
       method: 'POST',
       headers: {
-        'Content-Type': 'application/json'
+        'Content-Type': 'application/json',
+        ...headers
       },
       body: JSON.stringify({
         'text': commonStore.compositionParams.prompt.replaceAll(/<pad>|<start>|<end>/g, '').replaceAll('  ', ' ').trim()
@@ -175,7 +177,7 @@ const CompositionPanel: FC = observer(() => {
     });
   };
 
-  const onSubmit = (prompt: string) => {
+  const onSubmit = async (prompt: string) => {
     commonStore.setCompositionSubmittedPrompt(prompt);
 
     if (commonStore.status.status === ModelStatus.Offline && !commonStore.settings.apiUrl && commonStore.platform !== 'web') {
@@ -185,14 +187,22 @@ const CompositionPanel: FC = observer(() => {
     }
 
     let answer = '';
+    let finished = false;
+    const finish = () => {
+      finished = true;
+      commonStore.setCompositionGenerating(false);
+      generateNs(commonStore.compositionParams.autoPlay);
+    };
     compositionSseController = new AbortController();
+    const { url, headers } = await getReqUrl(port, '/v1/completions', true);
     fetchEventSource( // https://api.openai.com/v1/completions || http://127.0.0.1:${port}/v1/completions
-      getServerRoot(port, true) + '/v1/completions',
+      url,
       {
         method: 'POST',
         headers: {
           'Content-Type': 'application/json',
-          Authorization: `Bearer ${commonStore.settings.apiKey}`
+          Authorization: `Bearer ${commonStore.settings.apiKey}`,
+          ...headers
         },
         body: JSON.stringify({
           prompt,
@@ -205,9 +215,8 @@ const CompositionPanel: FC = observer(() => {
         signal: compositionSseController?.signal,
         onmessage(e) {
           scrollToBottom();
-          if (e.data.trim() === '[DONE]') {
+          if (!finished && e.data.trim() === '[DONE]') {
-            commonStore.setCompositionGenerating(false);
+            finish();
-            generateNs(commonStore.compositionParams.autoPlay);
             return;
           }
           let data;
@@ -220,13 +229,17 @@ const CompositionPanel: FC = observer(() => {
           if (data.model)
             commonStore.setLastModelName(data.model);
           if (data.choices && Array.isArray(data.choices) && data.choices.length > 0) {
+            if (!finished && data.choices[0]?.finish_reason) {
+              finish();
+              return;
+            }
             answer += data.choices[0]?.text || data.choices[0]?.delta?.content || '';
             setPrompt(prompt + answer.replace(/\s+$/, ''));
           }
         },
         async onopen(response) {
           if (response.status !== 200) {
-            toast(response.statusText + '\n' + (await response.text()), {
+            toast(response.status + ' - ' + response.statusText + ' - ' + (await response.text()), {
               type: 'error'
             });
           }
@@ -378,7 +391,7 @@ const CompositionPanel: FC = observer(() => {
               onSubmit(commonStore.compositionSubmittedPrompt);
             }} />
             <DialogButton className="grow" text={t('Reset')} title={t('Reset')}
-              contentText={t('Are you sure you want to reset this page? It cannot be undone.')}
+              content={t('Are you sure you want to reset this page? It cannot be undone.')}
               onConfirm={() => {
                 const isABC = commonStore.getCurrentModelConfig().modelParameters.modelName.toLowerCase().includes('abc');
                 const defaultPrompt = isABC ? defaultCompositionABCPrompt : defaultCompositionPrompt;

frontend/src/pages/Home.tsx

@@ -37,7 +37,7 @@ const clientNavCards: NavCard[] = [
   },
   {
     label: 'Configs',
-    desc: 'Manage your configs',
+    desc: 'Manage your configs, adjust the starting model and parameters',
     path: '/configs',
     icon: <DocumentSettings20Regular />
   },

frontend/src/pages/PresetsManager/MessagesEditor.tsx

@@ -3,7 +3,7 @@ import { DragDropContext, Draggable, Droppable, DropResult } from 'react-beautif
 import commonStore from '../../stores/commonStore';
 import { observer } from 'mobx-react-lite';
 import { v4 as uuid } from 'uuid';
-import { Button, Card, Dropdown, Option, Textarea } from '@fluentui/react-components';
+import { Card, Dropdown, Option, Textarea } from '@fluentui/react-components';
 import { useTranslation } from 'react-i18next';
 import { ToolTipButton } from '../../components/ToolTipButton';
 import { Delete20Regular, ReOrderDotsVertical20Regular } from '@fluentui/react-icons';
@@ -84,7 +84,10 @@ const MessagesEditor: FC = observer(() => {
 
   return (
     <div className="grid grid-cols-1 gap-2 overflow-hidden">
-      <Button style={{ width: '100%' }} onClick={createNewItem}>{t('New')}</Button>
+      <ToolTipButton text={t('New')}
+        desc={t('Create a new user or AI message content. You can prepare a chat record with AI here, and fill in the responses you want to get from AI in the tone of AI. When you use this preset, the chat record will be processed, and at this point, AI will better understand what you want it to do or what role to play.')}
+        style={{ width: '100%' }}
+        onClick={createNewItem} />
       <div className="overflow-x-hidden overflow-y-auto p-2">
         <DragDropContext onDragEnd={onDragEnd}>
           <Droppable droppableId="droppable">
@@ -123,7 +126,7 @@ const MessagesEditor: FC = observer(() => {
                             }}>
                             <Option value="user">{t('user')!}</Option>
                             <Option value="assistant">{t('assistant')!}</Option>
-                            {/* TODO <Option value="system">{t('system')!}</Option>*/}
+                            <Option value="system">{t('system')!}</Option>
                           </Dropdown>
                           <Textarea resize="vertical" className="grow" value={item.content}
                             style={{ minWidth: 0, borderRadius: 0 }}

frontend/src/pages/PresetsManager/PresetsButton.tsx

@@ -4,6 +4,7 @@ import React, { FC, lazy, PropsWithChildren, ReactElement, useState } from 'reac
 import {
   Button,
   Dialog,
+  DialogActions,
   DialogBody,
   DialogContent,
   DialogSurface,
@@ -60,10 +61,15 @@ const MessagesEditor = lazy(() => import('./MessagesEditor'));
 
 const PresetCardFrame: FC<PropsWithChildren & {
   onClick?: React.MouseEventHandler<HTMLButtonElement>
+  highlight?: boolean
 }> = (props) => {
   return <Button
     className="flex flex-col gap-1 w-32 h-56 break-all"
-    style={{ minWidth: 0, borderRadius: '0.75rem', justifyContent: 'unset' }}
+    style={{
+      minWidth: 0,
+      borderRadius: '0.75rem',
+      justifyContent: 'unset', ...(props.highlight ? { borderColor: '#115ea3', borderWidth: '2px' } : {})
+    }}
     onClick={props.onClick}
   >
     {props.children}
@@ -71,31 +77,28 @@ const PresetCardFrame: FC<PropsWithChildren & {
 };
 
 const PresetCard: FC<{
-  avatarImg: string,
-  name: string,
-  desc: string,
-  tag: string,
   editable: boolean,
+  preset: Preset,
   presetIndex: number,
-  onClick?: () => void
 }> = observer(({
-  avatarImg, name, desc, tag, editable, presetIndex, onClick
+  editable, preset, presetIndex
 }) => {
   const { t } = useTranslation();
 
-  return <PresetCardFrame onClick={(e) => {
+  return <PresetCardFrame highlight={commonStore.activePresetIndex === presetIndex}
-    if (onClick && e.currentTarget.contains(e.target as Node))
+    onClick={(e) => {
-      onClick();
+      if (e.currentTarget.contains(e.target as Node))
-  }}>
+        setActivePreset(presetIndex === -1 ? defaultPreset : preset, presetIndex);
-    <img src={absPathAsset(avatarImg)} className="rounded-xl select-none ml-auto mr-auto h-28" />
+    }}>
-    <Text size={400}>{name}</Text>
+    <img src={absPathAsset(preset.avatarImg)} className="rounded-xl select-none ml-auto mr-auto h-28" />
+    <Text size={400}>{preset.name}</Text>
     <Text size={200} style={{
       overflow: 'hidden', textOverflow: 'ellipsis',
       display: '-webkit-box', WebkitLineClamp: 3, WebkitBoxOrient: 'vertical'
-    }}>{desc}</Text>
+    }}>{preset.desc}</Text>
     <div className="grow" />
     <div className="flex justify-between w-full items-end">
-      <div className="text-xs font-thin text-gray-500">{t(tag)}</div>
+      <div className="text-xs font-thin text-gray-500">{t(preset.tag)}</div>
       {editable ?
         <ChatPresetEditor presetIndex={presetIndex} triggerButton={
           <ToolTipButton size="small" appearance="transparent" desc={t('Edit')} icon={<Edit20Regular />}
@@ -115,9 +118,11 @@ const ChatPresetEditor: FC<{
   presetIndex: number
 }> = observer(({ triggerButton, presetIndex }) => {
   const { t } = useTranslation();
+  const [open, setOpen] = React.useState(false);
+  const [showExitConfirm, setShowExitConfirm] = React.useState(false);
   const [editingMessages, setEditingMessages] = useState(false);
 
-  if (!commonStore.editingPreset)
+  if (open && !commonStore.editingPreset)
     commonStore.setEditingPreset({ ...defaultPreset });
   const editingPreset = commonStore.editingPreset!;
 
@@ -168,26 +173,46 @@ const ChatPresetEditor: FC<{
   };
 
   const savePreset = () => {
+    setOpen(false);
+    setShowExitConfirm(false);
     if (presetIndex === -1) {
       commonStore.setPresets([...commonStore.presets, { ...editingPreset }]);
-      setEditingPreset(defaultPreset);
     } else {
       commonStore.presets[presetIndex] = editingPreset;
       commonStore.setPresets(commonStore.presets);
     }
+    commonStore.setEditingPreset(null);
   };
 
   const activatePreset = () => {
     savePreset();
-    setActivePreset(editingPreset);
+    setActivePreset(editingPreset, presetIndex === -1 ? commonStore.presets.length - 1 : presetIndex);
   };
 
   const deletePreset = () => {
+    if (commonStore.activePresetIndex === presetIndex) {
+      setActivePreset(defaultPreset, -1);
+    }
     commonStore.presets.splice(presetIndex, 1);
     commonStore.setPresets(commonStore.presets);
+    setOpen(false);
+    setShowExitConfirm(false);
+    commonStore.setEditingPreset(null);
   };
 
-  return <Dialog>
+  return <Dialog open={open} onOpenChange={(e, data) => {
+    if (data.open) {
+      setOpen(true);
+    } else if (!commonStore.editingPreset ||
+      (presetIndex === -1 && JSON.stringify(editingPreset) === JSON.stringify(defaultPreset)) ||
+      (presetIndex !== -1 && JSON.stringify(editingPreset) === JSON.stringify(commonStore.presets[presetIndex]))) {
+      setOpen(false);
+      setShowExitConfirm(false);
+      commonStore.setEditingPreset(null);
+    } else {
+      setShowExitConfirm(true);
+    }
+  }}>
     <DialogTrigger disableButtonEnhancement>
       {triggerButton}
     </DialogTrigger>
@@ -196,17 +221,35 @@ const ChatPresetEditor: FC<{
       maxWidth: '80vw',
       maxHeight: '80vh',
       width: '500px',
-      height: '100%'
+      height: '100%',
+      transform: 'unset' // override the style for the new version of @fluentui/react-components to avoid conflicts with react-beautiful-dnd
     }}>
       <DialogBody style={{ height: '100%', overflow: 'hidden' }}>
-        <DialogContent className="flex flex-col gap-1 overflow-hidden">
+        {editingPreset && <DialogContent className="flex flex-col gap-1 overflow-hidden">
           <CustomToastContainer />
+          <Dialog open={showExitConfirm}>
+            <DialogSurface style={{ transform: 'unset' }}>
+              <DialogBody>
+                <DialogContent>
+                  {t('Content has been changed, are you sure you want to exit without saving?')}
+                </DialogContent>
+                <DialogActions>
+                  <Button appearance="secondary" onClick={() => {
+                    setShowExitConfirm(false);
+                  }}>{t('Cancel')}</Button>
+                  <Button appearance="primary" onClick={() => {
+                    setOpen(false);
+                    setShowExitConfirm(false);
+                    commonStore.setEditingPreset(null);
+                  }}>{t('Exit without saving')}</Button>
+                </DialogActions>
+              </DialogBody>
+            </DialogSurface>
+          </Dialog>
           <div className="flex justify-between">{
             presetIndex === -1
               ? <div />
-              : <DialogTrigger disableButtonEnhancement>
+              : <Button appearance="subtle" icon={<Delete20Regular />} onClick={deletePreset} />
-                <Button appearance="subtle" icon={<Delete20Regular />} onClick={deletePreset} />
-              </DialogTrigger>
           }
             <DialogTrigger disableButtonEnhancement>
               <Button appearance="subtle" icon={<Dismiss20Regular />} />
@@ -230,6 +273,7 @@ const ChatPresetEditor: FC<{
             editingMessages ?
               <div className="flex flex-col gap-1">
                 <Labeled flex spaceBetween label={t('Insert default system prompt at the beginning')}
+                  desc={t('Inside the model, there is a default prompt to improve the model\'s handling of common issues, but it may degrade the role-playing effect. You can disable this option to achieve a better role-playing effect.')}
                   content={
                     <Switch checked={editingPreset.presystem === undefined ? true : editingPreset.presystem}
                       onChange={(e, data) => {
@@ -239,6 +283,7 @@ const ChatPresetEditor: FC<{
                       }} />
                   } />
                 <Labeled flex breakline label={t('User Name')}
+                  desc={t('The name used internally by the model when processing user message, changing this value helps improve the role-playing effect.')}
                   content={
                     <Input placeholder="User" value={editingPreset.userName} onChange={(e, data) => {
                       setEditingPreset({
@@ -247,6 +292,7 @@ const ChatPresetEditor: FC<{
                     }} />
                   } />
                 <Labeled flex breakline label={t('Assistant Name')}
+                  desc={t('The name used internally by the model when processing AI message, changing this value helps improve the role-playing effect.')}
                   content={
                     <Input placeholder="Assistant" value={editingPreset.assistantName} onChange={(e, data) => {
                       setEditingPreset({
@@ -306,14 +352,10 @@ const ChatPresetEditor: FC<{
             <Button onClick={copyPreset}>{t('Copy')}</Button>
           </div>
           <div className="flex justify-between">
-            <DialogTrigger disableButtonEnhancement>
+            <Button appearance="primary" onClick={savePreset}>{t('Save')}</Button>
-              <Button appearance="primary" onClick={savePreset}>{t('Save')}</Button>
+            <Button appearance="primary" onClick={activatePreset}>{t('Activate')}</Button>
-            </DialogTrigger>
-            <DialogTrigger disableButtonEnhancement>
-              <Button appearance="primary" onClick={activatePreset}>{t('Activate')}</Button>
-            </DialogTrigger>
           </div>
-        </DialogContent>
+        </DialogContent>}
       </DialogBody>
     </DialogSurface>
   </Dialog>;
@@ -336,20 +378,16 @@ const ChatPresets: FC = observer(() => {
     {/*  </div>*/}
     {/*</PresetCardFrame>*/}
     <PresetCard
+      preset={defaultPreset}
       presetIndex={-1}
       editable={false}
-      onClick={() => {
-        setActivePreset(defaultPreset);
-      }} avatarImg={defaultPreset.avatarImg} name={defaultPreset.name} desc={defaultPreset.desc} tag={defaultPreset.tag}
     />
     {commonStore.presets.map((preset, index) => {
       return <PresetCard
+        key={index}
+        preset={preset}
         presetIndex={index}
         editable={true}
-        onClick={() => {
-          setActivePreset(preset);
-        }}
-        key={index} avatarImg={preset.avatarImg} name={preset.name} desc={preset.desc} tag={preset.tag}
       />;
     })}
   </div>;
@@ -418,7 +456,7 @@ export const PresetsButton: FC<{
       <ToolTipButton desc={t('Presets')} size={size} shape={shape} appearance={appearance}
         icon={<Accessibility28Regular />} />
     </DialogTrigger>
-    <DialogSurface style={{ paddingTop: 0, maxWidth: '90vw', width: 'fit-content' }}>
+    <DialogSurface style={{ paddingTop: 0, maxWidth: '90vw', width: 'fit-content', transform: 'unset' }}>
       <DialogBody>
         <DialogContent>
           <CustomToastContainer />

frontend/src/pages/Settings.tsx

@@ -17,6 +17,7 @@ import { useTranslation } from 'react-i18next';
 import { checkUpdate, toastWithButton } from '../utils';
 import { RestartApp } from '../../wailsjs/go/backend_golang/App';
 import { Language, Languages } from '../types/settings';
+import { toast } from 'react-toastify';
 
 export const GeneralSettings: FC = observer(() => {
   const { t } = useTranslation();
@@ -113,11 +114,15 @@ export const AdvancedGeneralSettings: FC = observer(() => {
                   commonStore.setSettings({
                     apiCompletionModelName: 'rwkv'
                   });
+              } else if (data.optionText === 'Ollama') {
+                toast(t('Don\'t forget to correctly fill in your Ollama API Chat Model Name.'),
+                  { type: 'info' });
               }
             }}>
             <Option value="">{t('Localhost')!}</Option>
             <Option value="https://rwkv.ai-creator.net/chntuned">RWKV</Option>
             <Option value="https://api.openai.com">OpenAI</Option>
+            <Option value="http://localhost:11434">Ollama</Option>
           </Dropdown>
         </div>
       } />
@@ -246,7 +251,7 @@ const Settings: FC = observer(() => {
                 }
                 {
                   commonStore.settings.language === 'zh' && commonStore.platform !== 'linux' &&
-                  <Labeled label={t('Use Tsinghua Pip Mirrors')} flex spaceBetween content={
+                  <Labeled label={t('Use Alibaba Cloud Pip Mirrors')} flex spaceBetween content={
                     <Switch checked={commonStore.settings.cnMirror}
                       onChange={(e, data) => {
                         commonStore.setSettings({

frontend/src/pages/Train.tsx

@@ -342,7 +342,7 @@ const LoraFinetune: FC = observer(() => {
       };
 
       if (msg === 'wsl is not enabled') {
-        enableWsl(false);
+        enableWsl(true);
       } else if (msg.includes('wsl.state: The system cannot find the file')) {
         enableWsl(true);
       } else {
@@ -391,7 +391,7 @@ const LoraFinetune: FC = observer(() => {
                     setDataParams({ dataPath: data.value });
                   }} />
                 <DialogButton text={t('Help')} title={t('Help')} markdown
-                  contentText={t('The data path should be a directory or a file in jsonl format (more formats will be supported in the future).\n\n' +
+                  content={t('The data path should be a directory or a file in jsonl format (more formats will be supported in the future).\n\n' +
                     'When you provide a directory path, all the txt files within that directory will be automatically converted into training data. ' +
                     'This is commonly used for large-scale training in writing, code generation, or knowledge bases.\n\n' +
                     'The jsonl format file can be referenced at https://github.com/josStorer/RWKV-Runner/blob/master/finetune/data/sample.jsonl.\n' +

frontend/src/startup.ts

@@ -9,6 +9,7 @@ import { t } from 'i18next';
 import { Preset } from './types/presets';
 import { toast } from 'react-toastify';
 import { MidiMessage, MidiPort } from './types/composition';
+import { throttle } from 'lodash-es';
 
 export async function startup() {
   initPresets();
@@ -16,6 +17,7 @@ export async function startup() {
   await GetPlatform().then(p => commonStore.setPlatform(p as Platform));
 
   if (commonStore.platform !== 'web') {
+    document.body.style.setProperty('overflow', 'hidden');
     downloadProgramFiles();
     EventsOn('downloadList', (data) => {
       if (data)
@@ -81,6 +83,7 @@ async function initConfig() {
   }).catch(() => {
     commonStore.setModelConfigs(commonStore.platform !== 'darwin' ? defaultModelConfigs : defaultModelConfigsMac, true);
   });
+  commonStore.setSettings({}, false); // to activate side effects
 }
 
 async function initCache(initUnfinishedModels: boolean) {
@@ -101,7 +104,7 @@ async function initPresets() {
 }
 
 async function initLoraModels() {
-  const refreshLoraModels = () => {
+  const refreshLoraModels = throttle(() => {
     ListDirFiles('lora-models').then((data) => {
       if (!data) return;
       const loraModels = [];
@@ -112,7 +115,7 @@ async function initLoraModels() {
       }
       commonStore.setLoraModels(loraModels);
     });
-  };
+  }, 2000);
 
   refreshLoraModels();
   EventsOn('fsnotify', (data: string) => {
@@ -122,9 +125,12 @@ async function initLoraModels() {
 }
 
 async function initLocalModelsNotify() {
+  const throttleRefreshLocalModels = throttle(() => {
+    refreshLocalModels({ models: commonStore.modelSourceList }, false); //TODO fix bug that only add models
+  }, 2000);
   EventsOn('fsnotify', (data: string) => {
     if (data.includes('models') && !data.includes('lora-models'))
-      refreshLocalModels({ models: commonStore.modelSourceList }, false); //TODO fix bug that only add models
+      throttleRefreshLocalModels();
   });
 }
 

frontend/src/stores/commonStore.ts

@@ -63,6 +63,7 @@ class CommonStore {
   monitorData: MonitorData | null = null;
   depComplete: boolean = false;
   platform: Platform = 'windows';
+  proxyPort: number = 0;
   lastModelName: string = '';
   // presets manager
   editingPreset: Preset | null = null;
@@ -74,6 +75,7 @@ class CommonStore {
   conversation: Conversation = {};
   conversationOrder: string[] = [];
   activePreset: Preset | null = null;
+  activePresetIndex: number = -1;
   attachmentUploading: boolean = false;
   attachments: {
     [uuid: string]: Attachment[]
@@ -259,13 +261,18 @@ class CommonStore {
   setSettings = (value: Partial<SettingsType>, saveConfig: boolean = true) => {
     this.settings = { ...this.settings, ...value };
 
-    if (this.settings.darkMode)
+    if (this.settings.darkMode) {
       WindowSetDarkTheme();
-    else
+      document.documentElement.setAttribute('style', 'color-scheme: dark;');
+    } else {
       WindowSetLightTheme();
+      document.documentElement.setAttribute('style', 'color-scheme: light;');
+    }
 
-    if (this.settings.language)
+    if (this.settings.language) {
       i18n.changeLanguage(this.settings.language);
+      document.documentElement.setAttribute('lang', this.settings.language === 'dev' ? 'en' : this.settings.language);
+    }
 
     if (saveConfig)
       saveConfigs();
@@ -317,6 +324,10 @@ class CommonStore {
     this.platform = value;
   }
 
+  setProxyPort(value: number) {
+    this.proxyPort = value;
+  }
+
   setCurrentInput(value: string) {
     this.currentInput = value;
   }
@@ -337,7 +348,7 @@ class CommonStore {
     this.lastUnfinishedModelDownloads = value;
   }
 
-  setEditingPreset(value: Preset) {
+  setEditingPreset(value: Preset | null) {
     this.editingPreset = value;
   }
 
@@ -351,6 +362,10 @@ class CommonStore {
     this.activePreset = value;
   }
 
+  setActivePresetIndex(value: number) {
+    this.activePresetIndex = value;
+  }
+
   setCompletionSubmittedPrompt(value: string) {
     this.completionSubmittedPrompt = value;
   }

frontend/src/style.scss

@@ -46,48 +46,24 @@ body {
   overflow-y: auto;
   overflow-x: hidden;
 
-  ul,
-  ol {
-    padding-left: 1.5em;
-  }
-
-  ol {
-    list-style: none;
-    counter-reset: item;
-
-    li {
-      counter-increment: item;
-
-      &::marker {
-        content: counter(item) '. ';
-      }
-    }
-  }
-
   pre {
     padding: 0;
+    background: transparent;
 
     code {
       font-size: 14px;
     }
   }
 
-  p {
-    margin: 0 0 10px;
-  }
-
   code {
-    padding: 0 0.4em;
-    margin: 0;
     white-space: pre-wrap;
     word-break: break-word;
     border-radius: 8px;
     background-color: var(--color-neutral-muted);
-    font-size: 11px;
+  }
 
-    .hljs {
+  details summary {
-      padding: 0;
+    cursor: pointer;
-    }
   }
 }
 

frontend/src/types/chat.ts

@@ -2,6 +2,7 @@ import { ApiParameters } from './configs';
 
 export const userName = 'M E';
 export const botName = 'A I';
+export const systemName = 'System';
 export const welcomeUuid = 'welcome';
 
 export enum MessageType {
@@ -9,7 +10,7 @@ export enum MessageType {
   Error
 }
 
-export type Side = 'left' | 'right'
+export type Side = 'left' | 'right' | 'center'
 export type Color = 'neutral' | 'brand' | 'colorful'
 export type MessageItem = {
   sender: string,

frontend/src/utils/index.tsx

@@ -2,6 +2,7 @@ import {
   AddToDownloadList,
   DeleteFile,
   DepCheck,
+  GetProxyPort,
   InstallPyDep,
   ListDirFiles,
   OpenOpenFileDialog,
@@ -26,9 +27,9 @@ import { DataProcessParameters, LoraFinetuneParameters } from '../types/train';
 import { InstrumentTypeNameMap, MidiMessage, tracksMinimalTotalTime } from '../types/composition';
 import logo from '../assets/images/logo.png';
 import { Preset } from '../types/presets';
-import { botName, Conversation, MessageType, Role, userName } from '../types/chat';
+import { botName, Conversation, MessageType, Role, systemName, userName } from '../types/chat';
 import { v4 as uuid } from 'uuid';
-import { findLastIndex } from 'lodash-es';
+import { findLastIndex, throttle } from 'lodash-es';
 
 export type Cache = {
   version: string
@@ -60,8 +61,7 @@ export async function refreshBuiltInModels(readCache: boolean = false) {
   else cache.models = manifest.models.slice();
 
   commonStore.setModelSourceList(cache.models);
-  await saveCache().catch(() => {
+  saveCache();
-  });
   return cache;
 }
 
@@ -119,8 +119,7 @@ export async function refreshLocalModels(cache: {
   commonStore.setModelSourceList(cache.models);
   if (initUnfinishedModels)
     initLastUnfinishedModelDownloads();
-  await saveCache().catch(() => {
+  saveCache();
-  });
 }
 
 function initLastUnfinishedModelDownloads() {
@@ -223,26 +222,34 @@ export const getStrategy = (modelConfig: ModelConfig | undefined = undefined) =>
   return strategy;
 };
 
-export const saveConfigs = async () => {
+export const saveConfigs = throttle(async () => {
-  const data: LocalConfig = {
+    const data: LocalConfig = {
-    modelSourceManifestList: commonStore.modelSourceManifestList,
+      modelSourceManifestList: commonStore.modelSourceManifestList,
-    currentModelConfigIndex: commonStore.currentModelConfigIndex,
+      currentModelConfigIndex: commonStore.currentModelConfigIndex,
-    modelConfigs: commonStore.modelConfigs,
+      modelConfigs: commonStore.modelConfigs,
-    settings: commonStore.settings,
+      settings: commonStore.settings,
-    dataProcessParams: commonStore.dataProcessParams,
+      dataProcessParams: commonStore.dataProcessParams,
-    loraFinetuneParams: commonStore.loraFinetuneParams
+      loraFinetuneParams: commonStore.loraFinetuneParams
-  };
+    };
-  return SaveJson('config.json', data);
+    return SaveJson('config.json', data);
-};
+  }, 500,
+  {
+    leading: true,
+    trailing: true
+  });
 
-export const saveCache = async () => {
+export const saveCache = throttle(async () => {
-  const data: Cache = {
+    const data: Cache = {
-    version: manifest.version,
+      version: manifest.version,
-    models: commonStore.modelSourceList,
+      models: commonStore.modelSourceList,
-    depComplete: commonStore.depComplete
+      depComplete: commonStore.depComplete
-  };
+    };
-  return SaveJson('cache.json', data);
+    return SaveJson('cache.json', data);
-};
+  }, 1000,
+  {
+    leading: true,
+    trailing: true
+  });
 
 export const savePresets = async () => {
   return SaveJson('presets.json', commonStore.presets);
@@ -310,6 +317,24 @@ export function bytesToReadable(size: number) {
   else return bytesToGb(size) + ' GB';
 }
 
+export async function getReqUrl(port: number, path: string, isCore: boolean = false): Promise<{
+  url: string,
+  headers: { [key: string]: string }
+}> {
+  const realUrl = getServerRoot(port, isCore) + path;
+  if (commonStore.platform === 'web')
+    return {
+      url: realUrl,
+      headers: {}
+    };
+  if (!commonStore.proxyPort)
+    await GetProxyPort().then(p => commonStore.setProxyPort(p));
+  return {
+    url: `http://127.0.0.1:${commonStore.proxyPort}`,
+    headers: { 'Real-Target': encodeURIComponent(realUrl) }
+  };
+}
+
 export function getServerRoot(defaultLocalPort: number, isCore: boolean = false) {
   const coreCustomApiUrl = commonStore.settings.coreApiUrl.trim().replace(/\/$/, '');
   if (isCore && coreCustomApiUrl)
@@ -581,8 +606,9 @@ export async function getSoundFont() {
   return soundUrl;
 }
 
-export const setActivePreset = (preset: Preset | null) => {
+export const setActivePreset = (preset: Preset | null, index: number) => {
   commonStore.setActivePreset(preset);
+  commonStore.setActivePresetIndex(index);
   //TODO if (preset.displayPresetMessages) {
   const { pushMessage, saveConversation } = newChatConversation();
   if (preset)
@@ -635,13 +661,13 @@ export function newChatConversation() {
     const newUuid = uuid();
     conversationOrder.push(newUuid);
     conversation[newUuid] = {
-      sender: role === 'user' ? userName : botName,
+      sender: role === 'user' ? userName : role === 'assistant' ? botName : systemName,
       type: MessageType.Normal,
-      color: role === 'user' ? 'brand' : 'colorful',
+      color: role === 'user' ? 'brand' : 'neutral',
       avatarImg: role === 'user' ? undefined : logo,
       time: new Date().toISOString(),
       content: content,
-      side: role === 'user' ? 'right' : 'left',
+      side: role === 'user' ? 'right' : role === 'assistant' ? 'left' : 'center',
       done: true
     };
   };

frontend/src/webWails.js

@@ -115,6 +115,12 @@ if (!window.go) {
   defineApp('ListDirFiles', async () => {
     return []
   })
+  defineApp('GetAbsPath', async (path) => {
+    return path
+  })
+  defineApp('GetProxyPort', async () => {
+    return 0
+  })
   defineApp('OpenOpenFileDialog', webOpenOpenFileDialog)
   defineApp('OpenSaveFileDialog', async (filterPattern, defaultFileName, savedContent) => {
     const saver = await import('file-saver')
@@ -127,7 +133,11 @@ if (!window.go) {
     return ''
   })
   defineApp('ReadJson', async (fileName) => {
-    return JSON.parse(localStorage.getItem(fileName))
+    const data = JSON.parse(localStorage.getItem(fileName))
+    if (data)
+      return data
+    else
+      throw new Error('File not found')
   })
   defineApp('SaveJson', async (fileName, data) => {
     localStorage.setItem(fileName, JSON.stringify(data))

frontend/tailwind.config.js

@@ -1,12 +1,121 @@
+const markdownElements = [
+  'div',
+  'p',
+  'span',
+
+  'video',
+  'img',
+
+  'abbr',
+  'acronym',
+  'b',
+  'blockquote',
+  'code',
+  'em',
+  'i',
+  'li',
+  'ol',
+  'ul',
+  'strong',
+  'table',
+  'tr',
+  'td',
+  'th',
+
+  'details',
+  'summary',
+  'kbd',
+  'samp',
+  'sub',
+  'sup',
+  'ins',
+  'del',
+  'var',
+  'q',
+  'dl',
+  'dt',
+  'dd',
+  'ruby',
+  'rt',
+  'rp',
+
+  'br',
+  'hr',
+
+  'h1',
+  'h2',
+  'h3',
+  'h4',
+  'h5',
+  'h6',
+
+  'thead',
+  'tbody',
+  'tfoot',
+  'u',
+  's',
+  'a',
+  'pre',
+  'cite'
+]
+
+const markdownPseudoElements = [
+  '::marker',
+  '::before',
+  '::after'
+]
+
+const tableElements = [
+  'table',
+  'tr',
+  'td',
+  'th',
+  'thead',
+  'tbody',
+  'tfoot'
+]
+
+const proseStyles = {
+  color: 'inherit'
+}
+
+const tableProseStyles = {
+  ...proseStyles,
+  borderWidth: 'thin',
+  borderColor: '#d2d2d5'
+}
+
+const elementsStyles = markdownElements.reduce((acc, element) => {
+  let styles = proseStyles
+  if (tableElements.includes(element))
+    styles = tableProseStyles
+
+  acc[element] = styles
+  markdownPseudoElements.forEach(pseudo => {
+    acc[element + pseudo] = styles
+  })
+  return acc
+}, {})
+
 /** @type {import('tailwindcss').Config} */
 export default {
   content: [
-    "./index.html",
+    './index.html',
-    "./src/**/*.{js,ts,jsx,tsx}",
+    './src/**/*.{js,ts,jsx,tsx}'
   ],
   theme: {
-    extend: {},
+    extend: {
+      typography: {
+        DEFAULT: {
+          css: {
+            color: 'inherit',
+            fontSize: 'inherit',
+            ...elementsStyles
+          }
+        }
+      }
+    }
   },
-  plugins: [],
+  plugins: [require('@tailwindcss/typography')]
 }
 

frontend/vite.config.ts

@@ -23,7 +23,7 @@ const embedded = [
   'react-beautiful-dnd',
   'react-draggable',
   '@magenta/music', 'html-midi-player',
-  'react-markdown', 'rehype-highlight', 'rehype-raw', 'remark-breaks', 'remark-gfm'
+  'react-markdown', 'rehype-highlight', 'rehype-raw', 'remark-breaks', 'remark-gfm', 'remark-math', 'rehype-katex', 'katex'
 ];
 
 function renderChunks(deps: Record<string, string>) {

frontend/wailsjs/go/backend_golang/App.d.ts

@@ -32,6 +32,8 @@ export function GetAbsPath(arg1:string):Promise<string>;
 
 export function GetPlatform():Promise<string>;
 
+export function GetProxyPort():Promise<number>;
+
 export function GetPyError():Promise<string>;
 
 export function InstallPyDep(arg1:string,arg2:boolean):Promise<string>;

frontend/wailsjs/go/backend_golang/App.js

@@ -62,6 +62,10 @@ export function GetPlatform() {
   return window['go']['backend_golang']['App']['GetPlatform']();
 }
 
+export function GetProxyPort() {
+  return window['go']['backend_golang']['App']['GetProxyPort']();
+}
+
 export function GetPyError() {
   return window['go']['backend_golang']['App']['GetPyError']();
 }

manifest.json

@@ -1,5 +1,5 @@
 {
-  "version": "1.7.3",
+  "version": "1.7.6",
   "introduction": {
     "en": "RWKV is an open-source, commercially usable large language model with high flexibility and great potential for development.\n### About This Tool\nThis tool aims to lower the barrier of entry for using large language models, making it accessible to everyone. It provides fully automated dependency and model management. You simply need to click and run, following the instructions, to deploy a local large language model. The tool itself is very compact and only requires a single executable file for one-click deployment.\nAdditionally, this tool offers an interface that is fully compatible with the OpenAI API. This means you can use any ChatGPT client as a client for RWKV, enabling capability expansion beyond just chat functionality.\n### Preset Configuration Rules at the Bottom\nThis tool comes with a series of preset configurations to reduce complexity. The naming rules for each configuration represent the following in order: device - required VRAM/memory - model size - model language.\nFor example, \"GPU-8G-3B-EN\" indicates that this configuration is for a graphics card with 8GB of VRAM, a model size of 3 billion parameters, and it uses an English language model.\nLarger model sizes have higher performance and VRAM requirements. Among configurations with the same model size, those with higher VRAM usage will have faster runtime.\nFor example, if you have 12GB of VRAM but running the \"GPU-12G-7B-EN\" configuration is slow, you can downgrade to \"GPU-8G-3B-EN\" for a significant speed improvement.\n### About RWKV\nRWKV is an RNN with Transformer-level LLM performance, which can also be directly trained like a GPT transformer (parallelizable). And it's 100% attention-free. You only need the hidden state at position t to compute the state at position t+1. You can use the \"GPT\" mode to quickly compute the hidden state for the \"RNN\" mode.<br/>So it's combining the best of RNN and transformer - great performance, fast inference, saves VRAM, fast training, \"infinite\" ctx_len, and free sentence embedding (using the final hidden state).",
     "zh": "RWKV是一个开源且允许商用的大语言模型，灵活性很高且极具发展潜力。\n### 关于本工具\n本工具旨在降低大语言模型的使用门槛，做到人人可用，本工具提供了全自动化的依赖和模型管理，你只需要直接点击运行，跟随引导，即可完成本地大语言模型的部署，工具本身体积极小，只需要一个exe即可完成一键部署。\n此外，本工具提供了与OpenAI API完全兼容的接口，这意味着你可以把任意ChatGPT客户端用作RWKV的客户端，实现能力拓展，而不局限于聊天。\n### 底部的预设配置规则\n本工具内置了一系列预设配置，以降低使用难度，每个配置名的规则，依次代表着：设备-所需显存/内存-模型规模-模型语言。\n例如，GPU-8G-3B-CN，表示该配置用于显卡，需要8G显存，模型规模为30亿参数，使用的是中文模型。\n模型规模越大，性能要求越高，显存要求也越高，而同样模型规模的配置中，显存占用越高的，运行速度越快。\n例如当你有12G显存，但运行GPU-12G-7B-CN配置速度比较慢，可降级成GPU-8G-3B-CN，将会大幅提速。\n### 关于RWKV\nRWKV是具有Transformer级别LLM性能的RNN，也可以像GPT Transformer一样直接进行训练（可并行化）。而且它是100% attention-free的。你只需在位置t处获得隐藏状态即可计算位置t + 1处的状态。你可以使用“GPT”模式快速计算用于“RNN”模式的隐藏状态。\n因此，它将RNN和Transformer的优点结合起来 - 高性能、快速推理、节省显存、快速训练、“无限”上下文长度以及免费的语句嵌入（使用最终隐藏状态）。"