Create rpc.py

rpc.py

@@ -0,0 +1,206 @@
+import tensorflow as tf
+from tensorflow import keras
+from keras.layers import *
+import keras_nlp
+
+import math
+import json
+from transformers import AutoTokenizer
+from tokenizers import AddedToken
+
+
+# Config
+input_size  = 512
+embed_dim   = 128
+
+
+# Tokenizer
+tokenizer = AutoTokenizer.from_pretrained('google/t5-v1_1-base')
+tokenizer.add_tokens(AddedToken("\n", normalized=False))
+tokenizer.add_tokens(AddedToken("<s>", normalized=False))
+vocab_size = len(tokenizer.get_vocab().keys())
+print("vocab_size:", vocab_size)
+print("pad token id:", tokenizer.pad_token)
+
+
+# Masked Accuracy Metric
+def masked_accuracy(y_true, y_pred, padding_token=tokenizer.pad_token_id):
+    y_true = tf.cast(y_true, tf.int32)
+    y_pred = tf.cast(tf.argmax(y_pred, axis=-1), tf.int32)
+    mask = tf.cast(tf.not_equal(y_true, padding_token), tf.float32)
+    matches = tf.cast(tf.equal(y_true, y_pred), tf.float32)
+    accuracy = tf.reduce_sum(matches * mask) / tf.reduce_sum(mask)
+    return accuracy
+
+
+# Embedding Layer
+class SharedEmbedding(tf.keras.layers.Layer):
+    def __init__(self, vocab_size, embed_dim, **kwargs):
+        super(SharedEmbedding, self).__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.embed_dim = embed_dim
+        
+    def build(self, input_shape):
+        self.shared_weights = self.add_weight(
+            shape=(self.vocab_size, self.embed_dim),
+            initializer='random_normal',
+            trainable=True,
+            name='shared_weights'
+        )
+        super(SharedEmbedding, self).build(input_shape)
+    
+    def call(self, inputs, mode='embedding', temp=0.1):
+        if mode == 'embedding':
+            return tf.nn.embedding_lookup(self.shared_weights, inputs)
+        elif mode == 'classify':
+            sw = tf.nn.l2_normalize(self.shared_weights, axis=-1)
+            return tf.nn.softmax(tf.matmul(inputs, sw, transpose_b=True)/temp, axis=-1)
+        
+
+# Attention Layer
+class Attention(keras.layers.Layer):
+    def __init__(self, **kwargs):
+        super(Attention, self).__init__(**kwargs)
+
+    def build(self, input_shape):
+        self.embed_dim = input_shape[-1]
+        self.mask = tf.where(tf.linalg.band_part(tf.ones((input_shape[-2], input_shape[-2])), -1, 0) == 1.0, 0.0, float("-inf"))
+        self.range_do = -tf.range(input_shape[-2])-1
+        self.range_undo = tf.range(input_shape[-2])+1
+        self.Q = self.add_weight(name='kernelQ',
+                                      shape=(input_shape[-1], input_shape[-1]),
+                                      initializer='uniform',
+                                      trainable=True)
+        self.K = self.add_weight(name='kernelK',
+                                      shape=(input_shape[-1], input_shape[-1]),
+                                      initializer='uniform',
+                                      trainable=True)
+        self.V = self.add_weight(name='kernelV',
+                                      shape=(input_shape[-1], input_shape[-1]),
+                                      initializer='uniform',
+                                      trainable=True)
+        super(Attention, self).build(input_shape)
+
+    def roll_embeddings(self, tensor, shift_values):
+        batch_size, time_size, embed_dim = tensor.shape
+        if batch_size is None: return tensor
+        shift_matrix   = tf.reshape(shift_values, (1, -1, 1))
+        shift_matrix   = tf.tile(shift_matrix, [batch_size, 1, embed_dim])
+        indices        = tf.range(embed_dim)
+        indices_matrix = tf.tile(indices, [batch_size * time_size])
+        indices_matrix = tf.reshape(indices_matrix, (batch_size, time_size, embed_dim))
+        new_indices    = (indices_matrix + shift_matrix) % embed_dim     
+        rolled_tensor  = tf.gather(tensor, new_indices, batch_dims=2)
+        return rolled_tensor
+
+    def call(self, x, pos):
+        q    = x @ self.Q
+        k    = x @ self.K
+        v    = x @ self.V
+        atti = tf.matmul(q, k,   transpose_b=True)
+        attp = tf.matmul(q, pos, transpose_b=True)
+        attp = self.roll_embeddings(attp, self.range_do)
+        att  = atti + attp
+        att  = tf.nn.softmax((att / math.sqrt(self.embed_dim)) + self.mask, axis=-1)
+        outi = att @ v
+        attp = self.roll_embeddings(att, self.range_undo)
+        outp = attp @ pos
+        out  = outi + outp
+        return out
+    
+
+# Encoder
+inputs = Input(shape=(input_size, ), dtype=tf.int32)
+emb_layer = SharedEmbedding(vocab_size, embed_dim)
+pos_layer = keras_nlp.layers.PositionEmbedding(input_size)
+
+x = LayerNormalization()(emb_layer(inputs, mode="embedding"))
+pos = pos_layer(x)
+
+b = 6
+for _ in range(b):
+    x += (2*b)**-0.5 * LayerNormalization()(Attention()(x, pos))
+    x += (2*b)**-0.5 * LayerNormalization()(Dense(embed_dim, activation="gelu")(x))
+x = tf.nn.l2_normalize(x, axis=-1)
+
+for _ in range(b):
+    x1 = Dense(embed_dim, activation="gelu")(x)
+    x1 = Dense(embed_dim, activation="gelu")(x1)
+    x += b**-0.5 * LayerNormalization()(x1)
+x = tf.nn.l2_normalize(x, axis=-1)
+
+x = emb_layer(x, mode="classify", temp=0.1)
+
+model = keras.Model(inputs=inputs, outputs=x)
+model.compile(
+    loss=keras.losses.SparseCategoricalCrossentropy(ignore_class=tokenizer.pad_token_id),
+    optimizer=keras.optimizers.AdamW(learning_rate=0.001),
+    metrics=[masked_accuracy, keras_nlp.metrics.Perplexity(mask_token_id=tokenizer.pad_token_id)],
+)
+
+
+# Import Model
+model.load_weights("rpc.keras")
+encoder = keras.Model(inputs=model.layers[0].input, outputs=model.layers[52].output)
+encoder.summary()
+
+
+# Vectorize Function
+def vectorize_texts(all_texts):
+    batch_size = 128
+    vects = []
+    for i in range(0, len(all_texts), batch_size):
+        texts = all_texts[i:i+batch_size]
+        toks = [text + ([tokenizer.pad_token_id] * (input_size - len(text))) for text in texts]
+        if len(toks) > 0:
+            toks = tf.constant(toks, shape=(len(toks), input_size))
+            vect = encoder.predict(toks, verbose=0)
+            for v, t in zip(vect, texts):
+                vects.append(v[:len(t), :])
+    return tf.concat(vects, axis=0).numpy()
+
+
+# Import Database and All Toks
+index = None
+all_toks = None
+def load_index(index_path="/dev/shm/rpc-vecdb/index"):
+    global index
+    global all_toks
+    #import ngtpy
+    #index = ngtpy.Index(index_path, read_only=True)
+    import faiss
+    index = faiss.read_index(index_path + "/index.faiss")
+    with open(index_path + "/all_toks.json", "r") as f:
+        all_toks = json.loads(f.read())
+
+
+# Generate Function
+def generate(text, use_rpc=True, max_tokens=128):
+    enc_text = tokenizer.encode(text, add_special_tokens=False)
+    
+    i = 0
+    while i < max_tokens and tok != vocab_size - 2:
+
+        enc_text = enc_text[-input_size:]
+        if use_rpc:
+            xq = vectorize_texts([enc_text])[-1]
+            #_id, _ = index.search(xq, size=1, epsilon=2)[0]
+            D, I = index.search(xq.reshape((1, -1)), 1)
+            _id = I[0][0]
+            if all_toks[_id] in carry_toks:
+                tmp = tf.argmax(tf.matmul(xq.reshape((1, -1)), encoder.layers[1].shared_weights, transpose_b=True), axis=-1).numpy()[0]
+                if all_toks[tmp] in enc_text: tok = tmp
+                else: tok = all_toks[_id]
+            else: tok = all_toks[_id]
+        else:
+            ins = enc_text + [tokenizer.pad_token_id] * (input_size - len(enc_text))
+            ins = tf.constant(ins, shape=(1, input_size))
+            res = model.predict(ins, verbose=0)[0][len(enc_text)-1]
+            tok = tf.argmax(res, axis=-1).numpy().tolist()
+        
+        enc_text += [tok]
+        new_text = tokenizer.decode(enc_text)
+        tok_str = new_text[len(text):]
+        text = new_text
+
+        yield tok_str