336 lines
13 KiB
Python
336 lines
13 KiB
Python
|
#!/usr/bin/env python
|
||
|
|
# pylint: disable=unused-argument
|
||
|
|
# This program is dedicated to the public domain under the CC0 license.
|
||
|
|
|
||
|
|
"""
|
||
|
|
Simple Bot to reply to Telegram messages.
|
||
|
|
|
||
|
|
First, a few handler functions are defined. Then, those functions are passed to
|
||
|
|
the Application and registered at their respective places.
|
||
|
|
Then, the bot is started and runs until we press Ctrl-C on the command line.
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
Basic Echobot example, repeats messages.
|
||
|
|
Press Ctrl-C on the command line or send a signal to the process to stop the
|
||
|
|
bot.
|
||
|
|
"""
|
||
|
|
import tokenmonster
|
||
|
|
tokenizer = tokenmonster.load("tokenizer_vocab/pol_vocab_exported")
|
||
|
|
|
||
|
|
import math
|
||
|
|
########################################################################################################
|
||
|
|
# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM
|
||
|
|
########################################################################################################
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
np.set_printoptions(precision=4, suppress=True, linewidth=200)
|
||
|
|
import types, torch, copy, time
|
||
|
|
from typing import List
|
||
|
|
torch.backends.cudnn.benchmark = True
|
||
|
|
torch.backends.cudnn.allow_tf32 = True
|
||
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
||
|
|
# torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = True
|
||
|
|
# torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = True
|
||
|
|
torch._C._jit_set_autocast_mode(False)
|
||
|
|
|
||
|
|
import torch.nn as nn
|
||
|
|
from torch.nn import functional as F
|
||
|
|
|
||
|
|
MyModule = torch.jit.ScriptModule
|
||
|
|
MyFunction = torch.jit.script_method
|
||
|
|
MyStatic = torch.jit.script
|
||
|
|
|
||
|
|
########################################################################################################
|
||
|
|
|
||
|
|
args = types.SimpleNamespace()
|
||
|
|
args.MODEL_NAME = 'tami-models/asshole-rwkv-6t'
|
||
|
|
args.n_layer = 12
|
||
|
|
args.n_embd = 768
|
||
|
|
args.vocab_size = 500
|
||
|
|
args.head_size = 64
|
||
|
|
|
||
|
|
#context = "מחפשת בדחיפות מחשב נייד הכי פשוט עם מטען"
|
||
|
|
|
||
|
|
|
||
|
|
########################################################################################################
|
||
|
|
|
||
|
|
class RWKV_RNN(MyModule):
|
||
|
|
def __init__(self, args):
|
||
|
|
super().__init__()
|
||
|
|
self.args = args
|
||
|
|
self.n_embd = args.n_embd
|
||
|
|
self.n_layer = args.n_layer
|
||
|
|
self.eval()
|
||
|
|
|
||
|
|
self.z = torch.load(args.MODEL_NAME + '.pth', map_location='cuda')
|
||
|
|
z = self.z
|
||
|
|
z['emb.weight'] = F.layer_norm(z['emb.weight'], (args.n_embd,), weight=z['blocks.0.ln0.weight'], bias=z['blocks.0.ln0.bias'])
|
||
|
|
|
||
|
|
keys = list(z.keys())
|
||
|
|
for k in keys:
|
||
|
|
if '.time_' in k: z[k] = z[k].squeeze()
|
||
|
|
if k.endswith('.time_decay'): z[k] = z[k].float()
|
||
|
|
if k.endswith('.time_faaaa'): z[k] = z[k].unsqueeze(-1).float()
|
||
|
|
for k in keys:
|
||
|
|
if k.endswith('maa_w'):
|
||
|
|
z[k.replace('maa_w','maa_wkvrg')] = torch.concat([z[k],z[k.replace('maa_w','maa_k')],z[k.replace('maa_w','maa_v')],z[k.replace('maa_w','maa_r')],z[k.replace('maa_w','maa_g')]]).clone().reshape(5, -1)
|
||
|
|
del z[k]
|
||
|
|
del z[k.replace('maa_w','maa_k')]
|
||
|
|
del z[k.replace('maa_w','maa_v')]
|
||
|
|
del z[k.replace('maa_w','maa_r')]
|
||
|
|
del z[k.replace('maa_w','maa_g')]
|
||
|
|
|
||
|
|
self.n_head = z['blocks.0.att.time_faaaa'].shape[0]
|
||
|
|
self.head_size = z['blocks.0.ln1.weight'].shape[0] // self.n_head
|
||
|
|
assert self.head_size == args.head_size
|
||
|
|
|
||
|
|
@MyFunction
|
||
|
|
def forward(self, token:int, state:List[torch.Tensor]):
|
||
|
|
with torch.no_grad():
|
||
|
|
z = self.z
|
||
|
|
x = z['emb.weight'][token]
|
||
|
|
for i in range(self.n_layer):
|
||
|
|
bbb = f'blocks.{i}.'
|
||
|
|
att = f'blocks.{i}.att.'
|
||
|
|
ffn = f'blocks.{i}.ffn.'
|
||
|
|
|
||
|
|
xx = F.layer_norm(x, (self.n_embd,), weight=z[bbb+'ln1.weight'], bias=z[bbb+'ln1.bias'])
|
||
|
|
xx, state[i*3+0], state[i*3+1] = time_mixing(self.n_head, self.head_size, xx, state[i*3+0], state[i*3+1],
|
||
|
|
z[att+'time_maa_x'], z[att+'time_maa_wkvrg'], z[att+'time_maa_w1'], z[att+'time_maa_w2'],
|
||
|
|
z[att+'time_decay_w1'], z[att+'time_decay_w2'], z[att+'time_faaaa'], z[att+'time_decay'],
|
||
|
|
z[att+'key.weight'], z[att+'value.weight'], z[att+'receptance.weight'], z[att+'gate.weight'], z[att+'output.weight'],
|
||
|
|
z[att+'ln_x.weight'], z[att+'ln_x.bias'])
|
||
|
|
x = x + xx
|
||
|
|
|
||
|
|
xx = F.layer_norm(x, (self.n_embd,), weight=z[bbb+'ln2.weight'], bias=z[bbb+'ln2.bias'])
|
||
|
|
xx, state[i*3+2] = channel_mixing(xx, state[i*3+2],
|
||
|
|
z[ffn+'time_maa_k'], z[ffn+'time_maa_r'],
|
||
|
|
z[ffn+'key.weight'], z[ffn+'value.weight'], z[ffn+'receptance.weight'])
|
||
|
|
x = x + xx
|
||
|
|
|
||
|
|
x = F.layer_norm(x, (self.n_embd,), weight=z['ln_out.weight'], bias=z['ln_out.bias'])
|
||
|
|
x = z['head.weight'] @ x
|
||
|
|
return x, state
|
||
|
|
|
||
|
|
########################################################################################################
|
||
|
|
|
||
|
|
def time_mixing__(H:int, N:int, x, x_prev, state, maa_x, maa_wkvrg, tm_w1, tm_w2, td_w1, td_w2, time_faaaa, time_decay, kw, vw, rw, gw, ow, ln_w, ln_b):
|
||
|
|
sx = x_prev - x
|
||
|
|
xxx = x + sx * maa_x # C
|
||
|
|
xxx = torch.tanh(xxx @ tm_w1).view(5, 1, -1) # C @ C*5L => 5L => 5*1*L
|
||
|
|
xxx = torch.bmm(xxx, tm_w2).view(5, -1) # 5*1*L @ 5*L*C => 5*1*C => 5*C
|
||
|
|
xxx = xxx + maa_wkvrg
|
||
|
|
xxx = xxx * sx.expand(5, -1) + x.expand(5, -1)
|
||
|
|
w, k, v, r, g = xxx.unbind(dim=0)
|
||
|
|
|
||
|
|
w = torch.tanh(w @ td_w1) @ td_w2
|
||
|
|
w = w.float() + time_decay
|
||
|
|
# assert w.dtype == torch.float
|
||
|
|
w = torch.exp(-torch.exp(w))
|
||
|
|
|
||
|
|
k = (kw @ k).view(H, N, 1)
|
||
|
|
v = (vw @ v).view(H, 1, N)
|
||
|
|
r = (rw @ r).view(H, 1, N)
|
||
|
|
g = torch.nn.functional.silu(gw @ g)
|
||
|
|
|
||
|
|
kv = (k @ v).float()
|
||
|
|
out = r @ (time_faaaa * kv + state).to(torch.bfloat16)
|
||
|
|
|
||
|
|
state = kv + w.view(H, N, 1) * state
|
||
|
|
|
||
|
|
out = torch.nn.functional.group_norm(out.view(1, H*N), num_groups=H, weight=ln_w, bias=ln_b, eps = 64e-5).view(H*N) # same as gn(x/8, eps=1e-5)
|
||
|
|
return ow @ (out * g), x, state
|
||
|
|
try:
|
||
|
|
time_mixing = torch.compile(time_mixing__, mode="max-autotune", fullgraph=True, dynamic=False)
|
||
|
|
except:
|
||
|
|
time_mixing = torch.jit.script(time_mixing__)
|
||
|
|
|
||
|
|
########################################################################################################
|
||
|
|
|
||
|
|
def channel_mixing__(x, x_prev, time_maa_k, time_maa_r, kw, vw, rw):
|
||
|
|
sx = x_prev - x
|
||
|
|
k = x + sx * time_maa_k
|
||
|
|
r = x + sx * time_maa_r
|
||
|
|
|
||
|
|
r = torch.sigmoid(rw @ r)
|
||
|
|
k = torch.relu(kw @ k) ** 2
|
||
|
|
|
||
|
|
return r * (vw @ k), x
|
||
|
|
try:
|
||
|
|
channel_mixing = torch.compile(channel_mixing__, mode="max-autotune", fullgraph=True, dynamic=False)
|
||
|
|
except:
|
||
|
|
channel_mixing = torch.jit.script(channel_mixing__)
|
||
|
|
|
||
|
|
########################################################################################################
|
||
|
|
|
||
|
|
print(f'\nUsing CUDA bf16. Loading {args.MODEL_NAME} ...')
|
||
|
|
model = RWKV_RNN(args)
|
||
|
|
|
||
|
|
TEMPERATURE = 1.5
|
||
|
|
NUM_TRIALS = 1
|
||
|
|
LENGTH_PER_TRIAL = 2048
|
||
|
|
MIN_P = 0.1
|
||
|
|
EOS_DECAY_DIV = 25
|
||
|
|
EOS_IGNORE = 25
|
||
|
|
|
||
|
|
print(f'\nPreprocessing context (slow version. see v2/rwkv/model.py for fast version)')
|
||
|
|
|
||
|
|
global_init_state = [None for _ in range(args.n_layer * 3)]
|
||
|
|
for i in range(args.n_layer): # state: 0=att_x_prev 1=att_kv 2=ffn_x_prev
|
||
|
|
global_init_state[i*3+0] = torch.zeros(args.n_embd, dtype=torch.bfloat16, requires_grad=False, device="cuda")
|
||
|
|
global_init_state[i*3+1] = torch.zeros((args.n_embd // args.head_size, args.head_size, args.head_size), dtype=torch.float, requires_grad=False, device="cuda")
|
||
|
|
global_init_state[i*3+2] = torch.zeros(args.n_embd, dtype=torch.bfloat16, requires_grad=False, device="cuda")
|
||
|
|
|
||
|
|
def generate_text(context_str):
|
||
|
|
context_tokens = tokenizer.tokenize(context_str)
|
||
|
|
print(context_tokens)
|
||
|
|
state = global_init_state.copy()
|
||
|
|
for token in context_tokens:
|
||
|
|
out, state = model.forward(token, state)
|
||
|
|
|
||
|
|
########################################################################################################
|
||
|
|
|
||
|
|
all_tokens = []
|
||
|
|
all_tokens += context_tokens.tolist()
|
||
|
|
|
||
|
|
added_tokens_count = 0
|
||
|
|
for i in range(LENGTH_PER_TRIAL):
|
||
|
|
added_tokens_count+=1
|
||
|
|
|
||
|
|
# applying temperature param
|
||
|
|
scores = out / TEMPERATURE
|
||
|
|
|
||
|
|
min_p = MIN_P
|
||
|
|
|
||
|
|
# convert logits to probabilities for min_p filtering
|
||
|
|
probs = torch.softmax(scores.float(),dim=-1)
|
||
|
|
|
||
|
|
# decaying eos token to induce longer generation
|
||
|
|
if (added_tokens_count <= EOS_IGNORE):
|
||
|
|
decay_factor = 0
|
||
|
|
else:
|
||
|
|
decay_factor = 1.0 - math.exp(-((added_tokens_count-EOS_IGNORE)/EOS_DECAY_DIV))
|
||
|
|
probs[0] *= decay_factor
|
||
|
|
|
||
|
|
# Get the probability of the top token for each sequence in the batch
|
||
|
|
max_prob, _ = (probs).max(dim=-1, keepdim=True)
|
||
|
|
# Calculate the actual min_p threshold by scaling min_p with the top token's probability
|
||
|
|
scaled_min_p = min_p * max_prob
|
||
|
|
# Create a mask for tokens that have a probability less than the scaled min_p
|
||
|
|
indices_to_remove = probs < scaled_min_p
|
||
|
|
|
||
|
|
scores.masked_fill_(indices_to_remove, float("-Inf"))
|
||
|
|
#print(scores[scores != float("-Inf")])
|
||
|
|
|
||
|
|
# convert min_p filtered logits to probabilities
|
||
|
|
probs = torch.softmax(scores.float(),dim=-1)
|
||
|
|
# reapply eos token decay to new probabilities
|
||
|
|
probs[0] *= decay_factor
|
||
|
|
|
||
|
|
token = torch.multinomial(probs, num_samples=1).squeeze(dim=-1)
|
||
|
|
|
||
|
|
#print(token)
|
||
|
|
|
||
|
|
all_tokens += [token]
|
||
|
|
if len(all_tokens) > 3:
|
||
|
|
# if the last 3 tokens are all the same break
|
||
|
|
if all_tokens[-1] == all_tokens[-2] == all_tokens[-3]: break
|
||
|
|
|
||
|
|
# find if any one token is overrepresented in the generated output and then assume the output is repeating
|
||
|
|
if len(all_tokens) > 256:
|
||
|
|
if any(num > len(all_tokens/4) for num in np.unique_counts(all_tokens).counts): break
|
||
|
|
|
||
|
|
if token == 0: break
|
||
|
|
|
||
|
|
out, state = model.forward(token, state)
|
||
|
|
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
|
||
|
|
# we generated all the tokens we could
|
||
|
|
output_string = tokenizer.decode(all_tokens)
|
||
|
|
logger.info(f"Generated {len(all_tokens)} = ctx:{len(context_tokens)} + gen:{added_tokens_count} tokens, Result:{output_string}")
|
||
|
|
|
||
|
|
# insert bolding of the input string TODO: wonky
|
||
|
|
#output_string = output_string[:len(context_str)] + "**" + output_string[len(context_str):]
|
||
|
|
|
||
|
|
return output_string
|
||
|
|
|
||
|
|
import logging
|
||
|
|
|
||
|
|
from telegram import ForceReply, Update
|
||
|
|
from telegram.ext import Application, CommandHandler, ContextTypes, MessageHandler, filters
|
||
|
|
# Read telegram token from .telegram_token
|
||
|
|
with open(".telegram_token") as f:
|
||
|
|
TOKEN = f.read().strip()
|
||
|
|
|
||
|
|
|
||
|
|
# Enable logging
|
||
|
|
logging.basicConfig(
|
||
|
|
filename='tamilm.log',
|
||
|
|
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO
|
||
|
|
)
|
||
|
|
# set higher logging level for httpx to avoid all GET and POST requests being logged
|
||
|
|
logging.getLogger("httpx").setLevel(logging.WARNING)
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
# Define a few command handlers. These usually take the two arguments update and
|
||
|
|
# context.
|
||
|
|
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||
|
|
"""Send a message when the command /start is issued."""
|
||
|
|
await update.message.reply_text("use /complete to complete your hebrew text in groupchats")
|
||
|
|
|
||
|
|
|
||
|
|
async def help_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||
|
|
"""Send a message when the command /help is issued."""
|
||
|
|
await update.message.reply_text("use /complete to complete your hebrew text")
|
||
|
|
|
||
|
|
async def complete_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||
|
|
"""Send a message when the command /help is issued."""
|
||
|
|
message_text = ' '.join(context.args)
|
||
|
|
|
||
|
|
# if this is a reply to a message include it in the prompt
|
||
|
|
if update.message.reply_to_message != None:
|
||
|
|
message_text = update.message.reply_to_message.text + ' ' + message_text
|
||
|
|
|
||
|
|
# format input into req structure
|
||
|
|
message_text = f"[REQ]\n{message_text}\n[RES]\n"
|
||
|
|
|
||
|
|
generated_responce = generate_text(message_text)
|
||
|
|
#logger.info(generated_responce)
|
||
|
|
await update.message.reply_text(generated_responce)
|
||
|
|
|
||
|
|
|
||
|
|
async def echo(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||
|
|
message_text = update.message.text
|
||
|
|
if update.message.chat.type == "private":
|
||
|
|
# format input into req structure
|
||
|
|
message_text = f"[REQ]\n{message_text}\n[RES]\n"
|
||
|
|
generated_responce = generate_text(message_text)
|
||
|
|
#(generated_responce)
|
||
|
|
await update.message.reply_text(generated_responce)
|
||
|
|
"""Echo the user message."""
|
||
|
|
#await update.message.reply_text(message_text)
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
"""Start the bot."""
|
||
|
|
# Create the Application and pass it your bot's token.
|
||
|
|
application = Application.builder().token(TOKEN).build()
|
||
|
|
|
||
|
|
# on different commands - answer in Telegram
|
||
|
|
application.add_handler(CommandHandler("start", start))
|
||
|
|
application.add_handler(CommandHandler("help", help_command))
|
||
|
|
application.add_handler(CommandHandler("complete", complete_command))
|
||
|
|
|
||
|
|
# on non command i.e message - echo the message on Telegram
|
||
|
|
application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, echo))
|
||
|
|
|
||
|
|
print("BOT READY")
|
||
|
|
|
||
|
|
# Run the bot until the user presses Ctrl-C
|
||
|
|
application.run_polling(allowed_updates=Update.ALL_TYPES)
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|