diff --git a/RWKV-v4neo/chat.py b/RWKV-v4neo/chat.py
index 4e95bd7..cb4a49e 100644
--- a/RWKV-v4neo/chat.py
+++ b/RWKV-v4neo/chat.py
@@ -140,6 +140,8 @@ def reply_msg(msg):
     print('Bot:', msg + '\n')
 
 def on_message(message):
+    global model_tokens, current_state
+
     srv = 'dummy_server'
 
     msg = message.strip()