Upload chat_final.py

chat_final.py

@@ -0,0 +1,142 @@
+print('Loading...')
+from src.model_run import RWKV_RNN
+import numpy as np
+import os, copy, types, gc, sys
+import torch
+from src.utils import TOKENIZER
+
+torch.backends.cudnn.benchmark = False
+torch.backends.cudnn.allow_tf32 = False
+torch.backends.cuda.matmul.allow_tf32 = False
+np.set_printoptions(precision=4, suppress=True, linewidth=200)
+
+WORD_NAME = ["20B_tokenizer.json", "20B_tokenizer.json"]
+UNKNOWN_CHAR = None
+tokenizer = TOKENIZER(WORD_NAME, UNKNOWN_CHAR=UNKNOWN_CHAR)
+
+args = types.SimpleNamespace()
+args.RUN_DEVICE = "cpu"  
+args.FLOAT_MODE = "fp32"
+args.vocab_size = 50277
+# args.MODEL_NAME = 'zrwkv-37fifth'
+args.MODEL_NAME = 'zrwkv-23fifth'
+
+args.n_layer = 12
+args.n_embd = 768
+args.ctx_len = 1024
+
+user = "User"
+bot = "Daniel"
+interface = ":"
+
+os.environ["RWKV_RUN_DEVICE"] = args.RUN_DEVICE
+MODEL_NAME = args.MODEL_NAME
+
+print(f'loading... {MODEL_NAME}')
+model = RWKV_RNN(args)
+
+model_tokens = []
+current_state = None
+
+def run_rnn(tokens, newline_adj = 0):
+    global model_tokens, current_state
+    for i in range(len(tokens)):
+        model_tokens += [int(tokens[i])]
+        if i == len(tokens) - 1:
+            out, current_state = model.forward(model_tokens, current_state)
+        else:
+            current_state = model.forward(model_tokens, current_state, preprocess_only = True)
+    
+    out[0] = -999999999  
+    out[187] += newline_adj
+    return out
+
+all_state = {}
+def save_all_stat(name, last_out):
+    all_state[name] = {}
+    all_state[name]['out'] = last_out
+    all_state[name]['rnn'] = copy.deepcopy(current_state)
+    all_state[name]['token'] = copy.deepcopy(model_tokens)
+
+def load_all_stat(name):
+    global model_tokens, current_state
+    current_state = copy.deepcopy(all_state[name]['rnn'])
+    model_tokens = copy.deepcopy(all_state[name]['token'])
+    return all_state[name]['out']
+
+print(f'\nRun prompt...')
+
+out = ""
+gc.collect()
+
+save_all_stat('chat_init', out)
+save_all_stat('chat', out)  # ensure that 'chat' key is added to all_state
+
+print(f'### prompt ###\n[{tokenizer.tokenizer.decode(model_tokens)}]\n')
+
+
+def reply_msg_generator():
+    while True:
+        msg = yield
+        print(f'{bot}{interface} {msg}\n')
+
+def on_message_generator():
+    global model_tokens, current_state
+    message = yield  # This yield allows us to receive the initial message
+    while True:
+        msg = message.replace('\\n','\n').strip()
+        if len(msg) > 10000:
+            message = yield 'your message is too long (max 1000 tokens)'
+
+        out = load_all_stat('chat')
+        new = f"{user}{interface} {msg}\n{bot}{interface}"
+        out = run_rnn(tokenizer.tokenizer.encode(new), newline_adj=-999999999)
+        save_all_stat('chat_pre', out)
+
+        begin = len(model_tokens)
+        out_last = begin
+        yield f'{bot}{interface}'  # Yield the bot's prompt immediately
+        for i in range(8000):
+            token = tokenizer.sample_logits(
+                out,
+                model_tokens,
+                args.ctx_len,
+                temperature=1.0,
+                top_p_usual=0.85,
+                top_p_newline=0.85,
+            )
+            out = run_rnn([token], newline_adj=1)
+
+            xxx = tokenizer.tokenizer.decode(model_tokens[out_last:])
+            if '\ufffd' not in xxx and 'user' not in str(xxx).lower() and '\n' not in xxx and str(xxx) != ':' and str(xxx) != '\n\n' and len(str(xxx)) > 0:
+                yield xxx  # Yield each part of the response as soon as it's ready
+                out_last = begin + i + 1
+            else:
+                out_last = begin + i + 1
+
+            send_msg = tokenizer.tokenizer.decode(model_tokens[begin:])
+            if '\ufffd' in send_msg or send_msg.endswith(f'{user}{interface}') or send_msg.endswith(f'{bot}{interface}') or '\n' in send_msg:
+                send_msg = send_msg.strip()
+                send_msg = send_msg.replace(f'{user}{interface}', '')
+                send_msg = send_msg.replace(f'{bot}{interface}', '')
+                send_msg = send_msg.replace('\n', '')
+                break   
+        save_all_stat('chat', out)
+        yield '\n'  # Yield a newline at the end of the response
+        message = yield  # Get the next message
+
+print('Start chatting with Daniel!')
+
+on_message_gen = on_message_generator()
+next_message = on_message_gen.__next__()  # Start the generator
+while True:
+    if next_message is None:  # If the generator is ready for a new message
+        msg = input(f'{user}{interface} ')
+        if len(msg.strip()) > 0:
+            next_message = on_message_gen.send(msg)  # Send the message to the generator and receive the next yield
+        else:
+            print('Error: please say something')
+    else:  # If the generator has yielded part of the response
+        print(next_message, end='', flush=True)
+        next_message = next(on_message_gen)  # Get the next part of the response
+