Implement basic chat mode

Add a chat mode. The functionalities are very basic, but hopefully work.

app.py

@@ -1,5 +1,9 @@
+from rwkv.utils import PIPELINE, PIPELINE_ARGS
+from rwkv.model import RWKV
 import gradio as gr
-import os, gc, torch
+import os
+import gc
+import torch
 from datetime import datetime
 from huggingface_hub import hf_hub_download
 from pynvml import *
@@ -11,31 +15,31 @@ desc = f'''Links:
 <a href='https://github.com/BlinkDL/ChatRWKV' target="_blank" style="margin:0 0.5em">ChatRWKV</a>
 <a href='https://github.com/BlinkDL/RWKV-LM' target="_blank" style="margin:0 0.5em">RWKV-LM</a>
 <a href="https://pypi.org/project/rwkv/" target="_blank" style="margin:0 0.5em">RWKV pip package</a>
-<a href="https://huggingface.co/spaces/BlinkDL/Raven-RWKV-7B" target="_blank" style="margin:0 0.5em">Raven 7B (alpaca-style)</a>
 '''
 
 os.environ["RWKV_JIT_ON"] = '1'
-os.environ["RWKV_CUDA_ON"] = '1' # if '1' then use CUDA kernel for seq mode (much faster)
+# if '1' then use CUDA kernel for seq mode (much faster)
+os.environ["RWKV_CUDA_ON"] = '1'
 
-from rwkv.model import RWKV
 model_path = hf_hub_download(repo_id="BlinkDL/rwkv-4-pile-14b", filename=f"{title}.pth")
-model = RWKV(model=model_path, strategy='cuda fp16i8 *24 -> cuda fp16')
-from rwkv.utils import PIPELINE, PIPELINE_ARGS
+model = RWKV(model=model_path, strategy='cuda fp16i8 *20 -> cuda fp16')
 pipeline = PIPELINE(model, "20B_tokenizer.json")
 
+########################################################################################################
+
 def infer(
         ctx,
         token_count=10,
         temperature=1.0,
         top_p=0.8,
-        presencePenalty = 0.1,
-        countPenalty = 0.1,
+        presence_enalty=0.1,
+        count_penalty=0.1,
 ):
-    args = PIPELINE_ARGS(temperature = max(0.2, float(temperature)), top_p = float(top_p),
-                     alpha_frequency = countPenalty,
-                     alpha_presence = presencePenalty,
-                     token_ban = [0], # ban the generation of some tokens
-                     token_stop = []) # stop generation whenever you see any token here
+    args = PIPELINE_ARGS(temperature=max(0.2, float(temperature)), top_p=float(top_p),
+                         alpha_frequency=float(count_penalty),
+                         alpha_presence=float(presence_enalty),
+                         token_ban=[0],  # ban the generation of some tokens
+                         token_stop=[])  # stop generation whenever you see any token here
 
     ctx = ctx.strip(' ')
     if ctx.endswith('\n'):
@@ -45,7 +49,7 @@ def infer(
 
     gpu_info = nvmlDeviceGetMemoryInfo(gpu_h)
     print(f'vram {gpu_info.total} used {gpu_info.used} free {gpu_info.free}')
-    
+
     all_tokens = []
     out_last = 0
     out_str = ''
@@ -66,7 +70,7 @@ def infer(
             occurrence[token] = 1
         else:
             occurrence[token] += 1
-        
+
         tmp = pipeline.decode(all_tokens[out_last:])
         if '\ufffd' not in tmp:
             out_str += tmp
@@ -106,10 +110,9 @@ Arrange the given numbers in ascending order.
     ["Simply put, the theory of relativity states that", 150, 1.0, 0.5, 0.2, 0.2],
 ]
 
-
-iface = gr.Interface(
+infer_interface = gr.Interface(
     fn=infer,
-    description=f'''{desc} *** <b>Please try examples first (bottom of page)</b> *** (edit them to use your question). Demo limited to ctxlen {ctx_limit}.''',
+    description=f'''{desc} <b>Please try examples first (bottom of page)</b> (edit them to use your question). Demo limited to ctxlen {ctx_limit}.''',
     allow_flagging="never",
     inputs=[
         gr.Textbox(lines=10, label="Prompt", value="Here's a short cyberpunk sci-fi adventure story. The story's main character is an artificial human created by a company called OpenBot.\n\nThe Story:\n"),  # prompt
@@ -124,10 +127,138 @@ iface = gr.Interface(
     cache_examples=False,
 ).queue()
 
+########################################################################################################
+
+user = "Bob"
+bot = "Alice"
+interface = ":"
+
+chat_intro = f'''
+The following is a coherent verbose detailed conversation between a girl named {bot} and her friend {user}. \
+{bot} is very intelligent, creative and friendly. \
+She is unlikely to disagree with {user}, and she doesn't like to ask {user} questions. \
+She also likes to tell {user} a lot about herself and her opinions, and she usually gives {user} kind, helpful and informative advices.
+
+{user}{interface} Hello, how are you doing?
+
+{bot}{interface} Hi {user}! Thanks, I'm fine. What about you?
+
+{user}{interface} I am fine. It's nice to see you. Look, here is a store selling tea and juice.
+
+{bot}{interface} Sure. Let's go inside. I would like to have some Mocha latte, which is my favourite!
+
+{user}{interface} What is it?
+
+{bot}{interface} Mocha latte is usually made with espresso, milk, chocolate, and frothed milk. Its flavors are frequently sweet.
+
+{user}{interface} Sounds tasty. I'll try it next time. Would you like to chat with me for a while?
+
+{bot}{interface} Of course! I'm glad to answer your questions or give helpful advices. You know, I am confident with my expertise. So please go ahead!
+
+'''
+
+_, intro_state = model.forward(pipeline.encode(chat_intro), None)
+
+def chat(
+        message: str,
+        history,
+        token_count=10,
+        temperature=1.0,
+        top_p=0.8,
+        presence_enalty=0.1,
+        count_penalty=0.1,
+):
+    args = PIPELINE_ARGS(temperature=max(0.2, float(temperature)), top_p=float(top_p),
+                         alpha_frequency=float(count_penalty),
+                         alpha_presence=float(presence_enalty),
+                         token_ban=[],  # ban the generation of some tokens
+                         token_stop=[])  # stop generation whenever you see any token here
+
+    message = message.strip(' ')
+    message = message.replace('\n', '')
+    ctx = f"{user}{interface} {message}\n\n{bot}{interface}"
+
+    gpu_info = nvmlDeviceGetMemoryInfo(gpu_h)
+    print(f'vram {gpu_info.total} used {gpu_info.used} free {gpu_info.free}')
+
+    history = history or [[], intro_state, []]  # [chat, state, all_tokens]
+
+    [chat_log, state, all_tokens] = history
+    out, state = model.forward(pipeline.encode(ctx)[-ctx_limit:], state)
+
+    begin = len(all_tokens)
+    out_last = begin
+    out_str: str = ''
+    occurrence = {}
+    for i in range(int(token_count)):
+        if i <= 0:
+            nl_bias = -float('inf')
+        elif i <= 30:
+            nl_bias = (i - 30) * 0.1
+        elif i <= 130:
+            nl_bias = 0
+        else:
+            nl_bias = (i - 130) * 0.25
+        out[187] += nl_bias
+        for n in occurrence:
+            out[n] -= (args.alpha_presence + occurrence[n] * args.alpha_frequency)
+
+        token = pipeline.sample_logits(out, temperature=args.temperature, top_p=args.top_p)
+        next_tokens = [token]
+        if token == 0:
+            next_tokens = pipeline.encode('\n\n')
+        all_tokens += next_tokens
+
+        if token not in occurrence:
+            occurrence[token] = 1
+        else:
+            occurrence[token] += 1
+
+        out, state = model.forward(next_tokens, state)
+
+        tmp = pipeline.decode(all_tokens[out_last:])
+        if '\ufffd' not in tmp:
+            print(tmp, end='', flush=True)
+            out_last = begin + i + 1
+
+        out_str = pipeline.decode(all_tokens[begin:])
+        out_str = out_str.replace("\r\n", '\n').replace('\\n', '\n')
+
+        if '\n\n' in out_str:
+            break
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    chat_log.append((message, out_str.strip()))
+    history = [chat_log, state, all_tokens]
+    return chat_log, history
+
+chat_interface = gr.Interface(
+    fn=chat,
+    description=f'''You are {user}, bot is {bot}.''',
+    allow_flagging="never",
+    inputs = [
+        gr.Textbox(label="Message"),
+        "state",
+        gr.Slider(10, 1000, step=10, value=250),    # token_count
+        gr.Slider(0.2, 2.0, step=0.1, value=1.0),   # temperature
+        gr.Slider(0.0, 1.0, step=0.05, value=0.8),  # top_p
+        gr.Slider(0.0, 1.0, step=0.1, value=0.2),   # presence_penalty
+        gr.Slider(0.0, 1.0, step=0.1, value=0.2),   # count_penalty
+    ],
+    outputs=[
+        gr.Chatbot(label="Chat Log", color_map=("blue", "pink")),
+        "state"
+    ]
+).queue()
+
+########################################################################################################
+
 demo = gr.TabbedInterface(
-    [iface], ["Generative"],
+    [infer_interface, chat_interface], ["Generative", "Chat"],
     title=title,
 )
 
 demo.queue(max_size=10)
-demo.launch(share=False)
+demo.launch(share=True)