pedrocas15/RPC-Chat

1import tensorflow as tf2from tensorflow import keras3from keras.layers import *4import keras_nlp5import subprocess6 7import math8import json9import spacy10from transformers import AutoTokenizer11from tokenizers import AddedToken12 13 14# Config15input_size  = 320#51216embed_dim   = 12817 18 19# Tokenizer20tokenizer = AutoTokenizer.from_pretrained('google/t5-v1_1-base')21tokenizer.add_tokens(AddedToken("\n", normalized=False))22tokenizer.add_tokens(AddedToken("<s>", normalized=False))23vocab_size = len(tokenizer.get_vocab().keys())24print("vocab_size:", vocab_size)25print("pad token id:", tokenizer.pad_token)26 27 28subprocess.run(["python", "-m", "spacy", "download", "en_core_web_lg"], check=True)29nlp = spacy.load("en_core_web_lg")30nlp.max_length = 200000031selected = {'NUM', 'PROPN'}32alltoks = sorted(list(tokenizer.get_vocab().items()), key=lambda x:x[1])33all_toks_text = "\n".join([t[0].replace("▁", "") for t in alltoks])34doc = nlp(all_toks_text)35carry_toks = set()36i = 037for ii, token in enumerate(doc):38    if str(token) in alltoks[i][0]: pass39    else: i += 140    if str(token) in alltoks[i][0] and token.pos_ in selected and i > 100:41        if (token.pos_ != "PROPN" or alltoks[i][0].replace("▁", "")[0].isupper()):42            carry_toks.add(alltoks[i][1])43print(len(carry_toks))44 45 46# Masked Accuracy Metric47def masked_accuracy(y_true, y_pred, padding_token=tokenizer.pad_token_id):48    y_true = tf.cast(y_true, tf.int32)49    y_pred = tf.cast(tf.argmax(y_pred, axis=-1), tf.int32)50    mask = tf.cast(tf.not_equal(y_true, padding_token), tf.float32)51    matches = tf.cast(tf.equal(y_true, y_pred), tf.float32)52    accuracy = tf.reduce_sum(matches * mask) / tf.reduce_sum(mask)53    return accuracy54 55 56# Embedding Layer57class SharedEmbedding(tf.keras.layers.Layer):58    def __init__(self, vocab_size, embed_dim, **kwargs):59        super(SharedEmbedding, self).__init__(**kwargs)60        self.vocab_size = vocab_size61        self.embed_dim = embed_dim62        63    def build(self, input_shape):64        self.shared_weights = self.add_weight(65            shape=(self.vocab_size, self.embed_dim),66            initializer='random_normal',67            trainable=True,68            name='shared_weights'69        )70        super(SharedEmbedding, self).build(input_shape)71    72    def call(self, inputs, mode='embedding', temp=0.1):73        if mode == 'embedding':74            return tf.nn.embedding_lookup(self.shared_weights, inputs)75        elif mode == 'classify':76            return tf.nn.softmax(tf.matmul(inputs, self.shared_weights, transpose_b=True), axis=-1) 77        78 79# Attention Layer80class DiffAttention(keras.layers.Layer):81    def __init__(self, depth, **kwargs):82        super(DiffAttention, self).__init__(**kwargs)83        self.lambda_init = 0.8 - 0.6 * math.exp(-0.3 * depth)84 85    def build(self, input_shape):86        self.embed_dim = input_shape[-1]87        self.input_size = input_shape[-2]88        self.mask = tf.where(tf.linalg.band_part(tf.ones((input_shape[-2], input_shape[-2])), -1, 0) == 1.0, 0.0, float("-inf"))89        self.range_do = -tf.range(input_shape[-2])-190        self.range_undo = tf.range(input_shape[-2])+191        self.Q = self.add_weight(name='kernelQ',92                                      shape=(input_shape[-1], input_shape[-1]),93                                      initializer='uniform',94                                      trainable=True)95        self.K = self.add_weight(name='kernelK',96                                      shape=(input_shape[-1], input_shape[-1]),97                                      initializer='uniform',98                                      trainable=True)99        self.V = self.add_weight(name='kernelV',100                                      shape=(input_shape[-1], input_shape[-1]),101                                      initializer='uniform',102                                      trainable=True)103 104        initializer = tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.1)105        self.lambda_q1 = self.add_weight(106            shape=(input_shape[-1],), initializer=initializer, trainable=True, name="lambda_q1"107        )108        self.lambda_k1 = self.add_weight(109            shape=(input_shape[-1],), initializer=initializer, trainable=True, name="lambda_k1"110        )111        self.lambda_q2 = self.add_weight(112            shape=(input_shape[-1],), initializer=initializer, trainable=True, name="lambda_q2"113        )114        self.lambda_k2 = self.add_weight(115            shape=(input_shape[-1],), initializer=initializer, trainable=True, name="lambda_k2"116        )117        118        super(DiffAttention, self).build(input_shape)119 120    def roll_embeddings(self, tensor, shift_values):121        batch_size, time_size, embed_dim = tensor.shape122        if batch_size is None: return tensor123        shift_matrix   = tf.reshape(shift_values, (1, -1, 1))124        shift_matrix   = tf.tile(shift_matrix, [batch_size, 1, embed_dim])125        indices        = tf.range(embed_dim)126        indices_matrix = tf.tile(indices, [batch_size * time_size])127        indices_matrix = tf.reshape(indices_matrix, (batch_size, time_size, embed_dim))128        new_indices    = (indices_matrix + shift_matrix) % embed_dim     129        rolled_tensor  = tf.gather(tensor, new_indices, batch_dims=2)130        return rolled_tensor131 132    def call(self, x, pos, pos_src):133        v    = x @ self.V134        q    = tf.transpose(tf.reshape(x @ self.Q, (-1, self.input_size, 2, self.embed_dim//2)), perm=[0, 2, 1, 3])135        k    = tf.transpose(tf.reshape(x @ self.K, (-1, self.input_size, 2, self.embed_dim//2)), perm=[0, 2, 1, 3])136        atti = tf.matmul(q, k,   transpose_b=True)137        attp = tf.matmul(q, pos, transpose_b=True)138        attp = self.roll_embeddings(tf.reshape(attp, (-1, self.input_size, self.input_size)), self.range_do)139        attp = tf.reshape(attp, (-1, 2, self.input_size, self.input_size))140        att  = atti + attp141        att  = tf.nn.softmax((att / math.sqrt(self.embed_dim)) + self.mask, axis=-1)142        att1 = att[:, 0]143        att2 = att[:, 1]144        145        # Differential attention146        lambda_1 = tf.math.exp(tf.reduce_sum(self.lambda_q1 * self.lambda_k1, axis=-1))147        lambda_2 = tf.math.exp(tf.reduce_sum(self.lambda_q2 * self.lambda_k2, axis=-1))148        lambda_full = lambda_1 - lambda_2 + self.lambda_init149        att = att1 - lambda_full * att2150 151        out = att @ v152        out = out * (1 - self.lambda_init)153        return out154    155 156# Import Model157model = keras.models.load_model(158    "rpc.keras",159    custom_objects={160        "DiffAttention" : DiffAttention,161        "SharedEmbedding" : SharedEmbedding,162        "masked_accuracy" : masked_accuracy163    }164)165encoder = keras.Model(inputs=model.layers[0].input, outputs=model.layers[-1].output)166encoder.summary()167 168 169# Vectorize Function170def vectorize_texts(all_texts):171    batch_size = 128172    vects = []173    for i in range(0, len(all_texts), batch_size):174        texts = all_texts[i:i+batch_size]175        toks = [text + ([tokenizer.pad_token_id] * (input_size - len(text))) for text in texts]176        if len(toks) > 0:177            toks = tf.constant(toks, shape=(len(toks), input_size))178            vect = encoder.predict(toks, verbose=0)179            for v, t in zip(vect, texts):180                vects.append(v[:len(t), :])181    return tf.concat(vects, axis=0).numpy()182 183 184# Import Database and All Toks185index = None186all_toks = None187index_type = None188def load_index(index_path="/dev/shm/rpc-vecdb/index", idx_type="ngt"):189    global index190    global all_toks191    global index_type192    index_type = idx_type193    if idx_type == "ngt":194        import ngtpy195        index = ngtpy.Index(index_path, read_only=True)196    elif idx_type == "faiss":    197        import faiss198        index = faiss.read_index(index_path + "/index.faiss")199    else:200        raise ValueError("Unknown index type")201    with open(index_path + "/all_toks.json", "r") as f:202        all_toks = json.loads(f.read())203 204 205# Generate Function206def generate(text, use_rpc=True, max_tokens=128):207    enc_text = tokenizer.encode(text, add_special_tokens=False)208    text = tokenizer.decode(enc_text)209    tok = None210    i = 0211    while i < max_tokens and tok != vocab_size - 2:212 213        enc_text = enc_text[-input_size:]214        if use_rpc:215            xq = vectorize_texts([enc_text])[-1]216            if index_type == "ngt":217                _id = index.search(xq, size=1, epsilon=1)[0][0]218            else:219                _id = index.search(xq.reshape((1, -1)), 1)[1][0][0]220            if all_toks[_id] in carry_toks:221                tmp = tf.argmax(tf.matmul(xq.reshape((1, -1)), encoder.layers[1].shared_weights, transpose_b=True), axis=-1).numpy()[0]222                if tmp in enc_text:223                    tok = tmp224                else: tok = all_toks[_id]225            else:226                tok = all_toks[_id]227        else:228            ins = enc_text + [tokenizer.pad_token_id] * (input_size - len(enc_text))229            ins = tf.constant(ins, shape=(1, input_size))230            res = model.predict(ins, verbose=0)[0][len(enc_text)-1]231            tok = tf.argmax(res, axis=-1).numpy().tolist()232        233        enc_text += [tok]234        new_text = tokenizer.decode(enc_text)235        res = new_text[len(text):]236        text = new_text237    238        yield res