Improve chat_with_bot.py script (#39)
This commit is contained in:
Alex
GitHub
parent
3587ff9e58
commit
c736ef5411
|
|
@ -1,11 +1,12 @@
|
|||
# Provides terminal-based chat interface for RWKV model.
|
||||
# Usage: python chat_with_bot.py C:\rwkv.cpp-169M.bin
|
||||
# Prompts and code adapted from https://github.com/BlinkDL/ChatRWKV/blob/9ca4cdba90efaee25cfec21a0bae72cbd48d8acd/chat.py
|
||||
|
||||
import os
|
||||
import sys
|
||||
import argparse
|
||||
import pathlib
|
||||
import copy
|
||||
from typing import List
|
||||
import torch
|
||||
import sampling
|
||||
import tokenizers
|
||||
import rwkv_cpp_model
|
||||
|
|
@ -13,89 +14,97 @@ import rwkv_cpp_shared_library
|
|||
|
||||
# ======================================== Script settings ========================================
|
||||
|
||||
# Copied from https://github.com/BlinkDL/ChatRWKV/blob/9ca4cdba90efaee25cfec21a0bae72cbd48d8acd/chat.py#L92-L178
|
||||
CHAT_LANG = 'English' # English // Chinese
|
||||
# English, Chinese
|
||||
LANGUAGE: str = 'English'
|
||||
|
||||
QA_PROMPT = False # True: Q & A prompt // False: chat prompt (need large model)
|
||||
# True: Q&A prompt
|
||||
# False: chat prompt (you need a large model for adequate quality, 7B+)
|
||||
QA_PROMPT: bool = False
|
||||
|
||||
if CHAT_LANG == 'English':
|
||||
interface = ':'
|
||||
MAX_GENERATION_LENGTH: int = 250
|
||||
|
||||
# Sampling temperature. It could be a good idea to increase temperature when top_p is low.
|
||||
TEMPERATURE: float = 0.8
|
||||
# For better Q&A accuracy and less diversity, reduce top_p (to 0.5, 0.2, 0.1 etc.)
|
||||
TOP_P: float = 0.5
|
||||
|
||||
if LANGUAGE == 'English':
|
||||
separator: str = ':'
|
||||
|
||||
if QA_PROMPT:
|
||||
user = "User"
|
||||
bot = "Bot" # Or: 'The following is a verbose and detailed Q & A conversation of factual information.'
|
||||
init_prompt = f'''
|
||||
The following is a verbose and detailed conversation between an AI assistant called {bot}, and a human user called {user}. {bot} is intelligent, knowledgeable, wise and polite.
|
||||
user: str = 'User'
|
||||
bot: str = 'Bot'
|
||||
init_prompt: str = f'''
|
||||
The following is a verbose and detailed conversation between an AI assistant called {bot}, and a human user called {user}. {bot} is intelligent, knowledgeable, wise and \
|
||||
polite.
|
||||
|
||||
{user}{interface} french revolution what year
|
||||
{user}{separator} french revolution what year
|
||||
|
||||
{bot}{interface} The French Revolution started in 1789, and lasted 10 years until 1799.
|
||||
{bot}{separator} The French Revolution started in 1789, and lasted 10 years until 1799.
|
||||
|
||||
{user}{interface} 3+5=?
|
||||
{user}{separator} 3+5=?
|
||||
|
||||
{bot}{interface} The answer is 8.
|
||||
{bot}{separator} The answer is 8.
|
||||
|
||||
{user}{interface} guess i marry who ?
|
||||
{user}{separator} guess i marry who ?
|
||||
|
||||
{bot}{interface} Only if you tell me more about yourself - what are your interests?
|
||||
{bot}{separator} Only if you tell me more about yourself - what are your interests?
|
||||
|
||||
{user}{interface} solve for a: 9-a=2
|
||||
{user}{separator} solve for a: 9-a=2
|
||||
|
||||
{bot}{interface} The answer is a = 7, because 9 - 7 = 2.
|
||||
{bot}{separator} The answer is a = 7, because 9 - 7 = 2.
|
||||
|
||||
{user}{interface} what is lhc
|
||||
{user}{separator} what is lhc
|
||||
|
||||
{bot}{interface} LHC is a high-energy particle collider, built by CERN, and completed in 2008. They used it to confirm the existence of the Higgs boson in 2012.
|
||||
{bot}{separator} LHC is a high-energy particle collider, built by CERN, and completed in 2008. They used it to confirm the existence of the Higgs boson in 2012.
|
||||
|
||||
'''
|
||||
else:
|
||||
user = "Bob"
|
||||
bot = "Alice"
|
||||
init_prompt = f'''
|
||||
user: str = 'Bob'
|
||||
bot: str = 'Alice'
|
||||
init_prompt: str = f'''
|
||||
The following is a verbose detailed conversation between {user} and a young girl {bot}. {bot} is intelligent, friendly and cute. {bot} is likely to agree with {user}.
|
||||
|
||||
{user}{interface} Hello {bot}, how are you doing?
|
||||
{user}{separator} Hello {bot}, how are you doing?
|
||||
|
||||
{bot}{interface} Hi {user}! Thanks, I'm fine. What about you?
|
||||
{bot}{separator} Hi {user}! Thanks, I'm fine. What about you?
|
||||
|
||||
{user}{interface} I am very good! It's nice to see you. Would you mind me chatting with you for a while?
|
||||
{user}{separator} I am very good! It's nice to see you. Would you mind me chatting with you for a while?
|
||||
|
||||
{bot}{interface} Not at all! I'm listening.
|
||||
{bot}{separator} Not at all! I'm listening.
|
||||
|
||||
'''
|
||||
|
||||
elif CHAT_LANG == 'Chinese':
|
||||
interface = ":"
|
||||
elif LANGUAGE == 'Chinese':
|
||||
separator: str = ':'
|
||||
|
||||
if QA_PROMPT:
|
||||
user = "Q"
|
||||
bot = "A"
|
||||
init_prompt = f'''
|
||||
user: str = 'Q'
|
||||
bot: str = 'A'
|
||||
init_prompt: str = f'''
|
||||
Expert Questions & Helpful Answers
|
||||
|
||||
Ask Research Experts
|
||||
|
||||
'''
|
||||
else:
|
||||
user = "Bob"
|
||||
bot = "Alice"
|
||||
init_prompt = f'''
|
||||
The following is a verbose and detailed conversation between an AI assistant called {bot}, and a human user called {user}. {bot} is intelligent, knowledgeable, wise and polite.
|
||||
user: str = 'Bob'
|
||||
bot: str = 'Alice'
|
||||
init_prompt: str = f'''
|
||||
The following is a verbose and detailed conversation between an AI assistant called {bot}, and a human user called {user}. {bot} is intelligent, knowledgeable, wise and \
|
||||
polite.
|
||||
|
||||
{user}{interface} what is lhc
|
||||
{user}{separator} what is lhc
|
||||
|
||||
{bot}{interface} LHC is a high-energy particle collider, built by CERN, and completed in 2008. They used it to confirm the existence of the Higgs boson in 2012.
|
||||
{bot}{separator} LHC is a high-energy particle collider, built by CERN, and completed in 2008. They used it to confirm the existence of the Higgs boson in 2012.
|
||||
|
||||
{user}{interface} 企鹅会飞吗
|
||||
{user}{separator} 企鹅会飞吗
|
||||
|
||||
{bot}{interface} 企鹅是不会飞的。它们的翅膀主要用于游泳和平衡，而不是飞行。
|
||||
{bot}{separator} 企鹅是不会飞的。它们的翅膀主要用于游泳和平衡，而不是飞行。
|
||||
|
||||
'''
|
||||
|
||||
FREE_GEN_LEN: int = 100
|
||||
|
||||
# Sampling settings.
|
||||
GEN_TEMP: float = 0.8 # It could be a good idea to increase temp when top_p is low
|
||||
GEN_TOP_P: float = 0.5 # Reduce top_p (to 0.5, 0.2, 0.1 etc.) for better Q&A accuracy (and less diversity)
|
||||
else:
|
||||
assert False, f'Invalid language {LANGUAGE}'
|
||||
|
||||
# =================================================================================================
|
||||
|
||||
|
|
@ -117,79 +126,87 @@ model = rwkv_cpp_model.RWKVModel(library, args.model_path)
|
|||
|
||||
prompt_tokens = tokenizer.encode(init_prompt).ids
|
||||
prompt_token_count = len(prompt_tokens)
|
||||
print(f'Processing {prompt_token_count} prompt tokens, may take a while')
|
||||
|
||||
|
||||
########################################################################################################
|
||||
|
||||
def run_rnn(tokens: List[int]):
|
||||
model_tokens: list[int] = []
|
||||
|
||||
logits, model_state = None, None
|
||||
|
||||
def process_tokens(_tokens: list[int]) -> torch.Tensor:
|
||||
global model_tokens, model_state, logits
|
||||
|
||||
tokens = [int(x) for x in tokens]
|
||||
model_tokens += tokens
|
||||
_tokens = [int(x) for x in _tokens]
|
||||
|
||||
for token in tokens:
|
||||
logits, model_state = model.eval(token, model_state, model_state, logits)
|
||||
model_tokens += _tokens
|
||||
|
||||
for _token in _tokens:
|
||||
logits, model_state = model.eval(_token, model_state, model_state, logits)
|
||||
|
||||
return logits
|
||||
|
||||
all_state = {}
|
||||
state_by_thread: dict[str, dict] = {}
|
||||
|
||||
def save_all_stat(thread: str, last_out):
|
||||
n = f'{thread}'
|
||||
all_state[n] = {}
|
||||
all_state[n]['logits'] = copy.deepcopy(last_out)
|
||||
all_state[n]['rnn'] = copy.deepcopy(model_state)
|
||||
all_state[n]['token'] = copy.deepcopy(model_tokens)
|
||||
def save_thread_state(_thread: str, _logits: torch.Tensor) -> None:
|
||||
state_by_thread[_thread] = {}
|
||||
state_by_thread[_thread]['logits'] = copy.deepcopy(_logits)
|
||||
state_by_thread[_thread]['rnn'] = copy.deepcopy(model_state)
|
||||
state_by_thread[_thread]['token'] = copy.deepcopy(model_tokens)
|
||||
|
||||
def load_all_stat(thread: str):
|
||||
def load_thread_state(_thread: str) -> torch.Tensor:
|
||||
global model_tokens, model_state
|
||||
n = f'{thread}'
|
||||
model_state = copy.deepcopy(all_state[n]['rnn'])
|
||||
model_tokens = copy.deepcopy(all_state[n]['token'])
|
||||
return copy.deepcopy(all_state[n]['logits'])
|
||||
model_state = copy.deepcopy(state_by_thread[_thread]['rnn'])
|
||||
model_tokens = copy.deepcopy(state_by_thread[_thread]['token'])
|
||||
return copy.deepcopy(state_by_thread[_thread]['logits'])
|
||||
|
||||
########################################################################################################
|
||||
|
||||
model_tokens = []
|
||||
logits, model_state = None, None
|
||||
print(f'Processing {prompt_token_count} prompt tokens, may take a while')
|
||||
|
||||
for token in prompt_tokens:
|
||||
logits, model_state = model.eval(token, model_state, model_state, logits)
|
||||
|
||||
model_tokens.append(token)
|
||||
|
||||
save_all_stat('chat_init', logits)
|
||||
print('\nChat initialized! Write something and press Enter.')
|
||||
save_all_stat('chat', logits)
|
||||
save_thread_state('chat_init', logits)
|
||||
save_thread_state('chat', logits)
|
||||
|
||||
print(f'\nChat initialized! Your name is {user}. Write something and press Enter. Use \\n to add line breaks to your message.')
|
||||
|
||||
while True:
|
||||
# Read user input
|
||||
user_input = input(f'> {user}{interface} ')
|
||||
msg = user_input.replace('\\n','\n').strip()
|
||||
user_input = input(f'> {user}{separator} ')
|
||||
msg = user_input.replace('\\n', '\n').strip()
|
||||
|
||||
temperature = TEMPERATURE
|
||||
top_p = TOP_P
|
||||
|
||||
if "-temp=" in msg:
|
||||
temperature = float(msg.split('-temp=')[1].split(' ')[0])
|
||||
|
||||
msg = msg.replace('-temp='+f'{temperature:g}', '')
|
||||
|
||||
temperature = GEN_TEMP
|
||||
top_p = GEN_TOP_P
|
||||
if ("-temp=" in msg):
|
||||
temperature = float(msg.split("-temp=")[1].split(" ")[0])
|
||||
msg = msg.replace("-temp="+f'{temperature:g}', "")
|
||||
# print(f"temp: {temperature}")
|
||||
if ("-top_p=" in msg):
|
||||
top_p = float(msg.split("-top_p=")[1].split(" ")[0])
|
||||
msg = msg.replace("-top_p="+f'{top_p:g}', "")
|
||||
# print(f"top_p: {top_p}")
|
||||
if temperature <= 0.2:
|
||||
temperature = 0.2
|
||||
|
||||
if temperature >= 5:
|
||||
temperature = 5
|
||||
|
||||
if "-top_p=" in msg:
|
||||
top_p = float(msg.split('-top_p=')[1].split(' ')[0])
|
||||
|
||||
msg = msg.replace('-top_p='+f'{top_p:g}', '')
|
||||
|
||||
if top_p <= 0:
|
||||
top_p = 0
|
||||
|
||||
msg = msg.strip()
|
||||
|
||||
# + reset --> reset chat
|
||||
if msg == '+reset':
|
||||
logits = load_all_stat('chat_init')
|
||||
save_all_stat('chat', logits)
|
||||
print(f'{bot}{interface} "Chat reset."\n')
|
||||
logits = load_thread_state('chat_init')
|
||||
save_thread_state('chat', logits)
|
||||
print(f'{bot}{separator} Chat reset.\n')
|
||||
continue
|
||||
elif msg[:5].lower() == '+gen ' or msg[:3].lower() == '+i ' or msg[:4].lower() == '+qa ' or msg[:4].lower() == '+qq ' or msg.lower() == '+++' or msg.lower() == '++':
|
||||
|
||||
|
|
@ -199,8 +216,8 @@ while True:
|
|||
# print(f'### prompt ###\n[{new}]')
|
||||
model_state = None
|
||||
model_tokens = []
|
||||
logits = run_rnn(tokenizer.encode(new).ids)
|
||||
save_all_stat('gen_0', logits)
|
||||
logits = process_tokens(tokenizer.encode(new).ids)
|
||||
save_thread_state('gen_0', logits)
|
||||
|
||||
# +i YOUR INSTRUCT --> free single-round generation with any instruct. Requires Raven model.
|
||||
elif msg[:3].lower() == '+i ':
|
||||
|
|
@ -215,8 +232,8 @@ Below is an instruction that describes a task. Write a response that appropriate
|
|||
# print(f'### prompt ###\n[{new}]')
|
||||
model_state = None
|
||||
model_tokens = []
|
||||
logits = run_rnn(tokenizer.encode(new).ids)
|
||||
save_all_stat('gen_0', logits)
|
||||
logits = process_tokens(tokenizer.encode(new).ids)
|
||||
save_thread_state('gen_0', logits)
|
||||
|
||||
# +qq YOUR QUESTION --> answer an independent question with more creativity (regardless of context).
|
||||
elif msg[:4].lower() == '+qq ':
|
||||
|
|
@ -224,25 +241,25 @@ Below is an instruction that describes a task. Write a response that appropriate
|
|||
# print(f'### prompt ###\n[{new}]')
|
||||
model_state = None
|
||||
model_tokens = []
|
||||
logits = run_rnn(tokenizer.encode(new).ids)
|
||||
save_all_stat('gen_0', logits)
|
||||
logits = process_tokens(tokenizer.encode(new).ids)
|
||||
save_thread_state('gen_0', logits)
|
||||
|
||||
# +qa YOUR QUESTION --> answer an independent question (regardless of context).
|
||||
elif msg[:4].lower() == '+qa ':
|
||||
logits = load_all_stat('chat_init')
|
||||
logits = load_thread_state('chat_init')
|
||||
|
||||
real_msg = msg[4:].strip()
|
||||
new = f"{user}{interface} {real_msg}\n\n{bot}{interface}"
|
||||
new = f"{user}{separator} {real_msg}\n\n{bot}{separator}"
|
||||
# print(f'### qa ###\n[{new}]')
|
||||
|
||||
logits = run_rnn(tokenizer.encode(new).ids)
|
||||
save_all_stat('gen_0', logits)
|
||||
logits = process_tokens(tokenizer.encode(new).ids)
|
||||
save_thread_state('gen_0', logits)
|
||||
|
||||
# +++ --> continue last free generation (only for +gen / +i)
|
||||
elif msg.lower() == '+++':
|
||||
try:
|
||||
logits = load_all_stat('gen_1')
|
||||
save_all_stat('gen_0', logits)
|
||||
logits = load_thread_state('gen_1')
|
||||
save_thread_state('gen_0', logits)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
continue
|
||||
|
|
@ -250,7 +267,7 @@ Below is an instruction that describes a task. Write a response that appropriate
|
|||
# ++ --> retry last free generation (only for +gen / +i)
|
||||
elif msg.lower() == '++':
|
||||
try:
|
||||
logits = load_all_stat('gen_0')
|
||||
logits = load_thread_state('gen_0')
|
||||
except Exception as e:
|
||||
print(e)
|
||||
continue
|
||||
|
|
@ -260,41 +277,46 @@ Below is an instruction that describes a task. Write a response that appropriate
|
|||
# + --> alternate chat reply
|
||||
if msg.lower() == '+':
|
||||
try:
|
||||
logits = load_all_stat('chat_pre')
|
||||
logits = load_thread_state('chat_pre')
|
||||
except Exception as e:
|
||||
print(e)
|
||||
continue
|
||||
# chat with bot
|
||||
else:
|
||||
logits = load_all_stat('chat')
|
||||
new = f"{user}{interface} {msg}\n\n{bot}{interface}"
|
||||
logits = load_thread_state('chat')
|
||||
new = f"{user}{separator} {msg}\n\n{bot}{separator}"
|
||||
# print(f'### add ###\n[{new}]')
|
||||
logits = run_rnn(tokenizer.encode(new).ids)
|
||||
save_all_stat('chat_pre', logits)
|
||||
logits = process_tokens(tokenizer.encode(new).ids)
|
||||
save_thread_state('chat_pre', logits)
|
||||
|
||||
thread = 'chat'
|
||||
|
||||
# Print bot response
|
||||
print(f"> {bot}{interface}", end='')
|
||||
print(f"> {bot}{separator}", end='')
|
||||
|
||||
decoded = ''
|
||||
begin = len(model_tokens)
|
||||
out_last = begin
|
||||
start_index: int = len(model_tokens)
|
||||
accumulated_tokens: list[int] = []
|
||||
|
||||
for i in range(FREE_GEN_LEN):
|
||||
token = sampling.sample_logits(logits, temperature, top_p)
|
||||
logits = run_rnn([token])
|
||||
decoded = tokenizer.decode(model_tokens[out_last:])
|
||||
if '\ufffd' not in decoded: # avoid utf-8 display issues
|
||||
for i in range(MAX_GENERATION_LENGTH):
|
||||
token: int = sampling.sample_logits(logits, temperature, top_p)
|
||||
|
||||
logits: torch.Tensor = process_tokens([token])
|
||||
|
||||
# Avoid UTF-8 display issues
|
||||
accumulated_tokens += [token]
|
||||
|
||||
decoded: str = tokenizer.decode(accumulated_tokens)
|
||||
|
||||
if '\uFFFD' not in decoded:
|
||||
print(decoded, end='', flush=True)
|
||||
out_last = begin + i + 1
|
||||
|
||||
accumulated_tokens = []
|
||||
|
||||
if thread == 'chat':
|
||||
send_msg = tokenizer.decode(model_tokens[begin:])
|
||||
if '\n\n' in send_msg:
|
||||
send_msg = send_msg.strip()
|
||||
if '\n\n' in tokenizer.decode(model_tokens[start_index:]):
|
||||
break
|
||||
if i == FREE_GEN_LEN - 1:
|
||||
|
||||
if i == MAX_GENERATION_LENGTH - 1:
|
||||
print()
|
||||
|
||||
save_all_stat(thread, logits)
|
||||
save_thread_state(thread, logits)
|
||||
|
|
|
|||
Loading…
Reference in New Issue