Update rpc.py

rpc.py

@@ -53,17 +53,18 @@ class SharedEmbedding(tf.keras.layers.Layer):
         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)
+            return tf.nn.softmax(tf.matmul(inputs, self.shared_weights, transpose_b=True), axis=-1) 
         
 
 # Attention Layer
-class Attention(keras.layers.Layer):
-    def __init__(self, **kwargs):
-        super(Attention, self).__init__(**kwargs)
+class DiffAttention(keras.layers.Layer):
+    def __init__(self, depth, **kwargs):
+        super(DiffAttention, self).__init__(**kwargs)
+        self.lambda_init = 0.8 - 0.6 * math.exp(-0.3 * depth)
 
     def build(self, input_shape):
         self.embed_dim = input_shape[-1]
+        self.input_size = input_shape[-2]
         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
@@ -79,7 +80,22 @@ class Attention(keras.layers.Layer):
                                       shape=(input_shape[-1], input_shape[-1]),
                                       initializer='uniform',
                                       trainable=True)
-        super(Attention, self).build(input_shape)
+
+        initializer = tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.1)
+        self.lambda_q1 = self.add_weight(
+            shape=(input_shape[-1],), initializer=initializer, trainable=True, name="lambda_q1"
+        )
+        self.lambda_k1 = self.add_weight(
+            shape=(input_shape[-1],), initializer=initializer, trainable=True, name="lambda_k1"
+        )
+        self.lambda_q2 = self.add_weight(
+            shape=(input_shape[-1],), initializer=initializer, trainable=True, name="lambda_q2"
+        )
+        self.lambda_k2 = self.add_weight(
+            shape=(input_shape[-1],), initializer=initializer, trainable=True, name="lambda_k2"
+        )
+        
+        super(DiffAttention, self).build(input_shape)
 
     def roll_embeddings(self, tensor, shift_values):
         batch_size, time_size, embed_dim = tensor.shape
@@ -93,55 +109,43 @@ class Attention(keras.layers.Layer):
         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
+    def call(self, x, pos, pos_src):
         v    = x @ self.V
+        q    = tf.transpose(tf.reshape(x @ self.Q, (-1, self.input_size, 2, self.embed_dim//2)), perm=[0, 2, 1, 3])
+        k    = tf.transpose(tf.reshape(x @ self.K, (-1, self.input_size, 2, self.embed_dim//2)), perm=[0, 2, 1, 3])
         atti = tf.matmul(q, k,   transpose_b=True)
         attp = tf.matmul(q, pos, transpose_b=True)
-        attp = self.roll_embeddings(attp, self.range_do)
+        attp = self.roll_embeddings(tf.reshape(attp, (-1, self.input_size, self.input_size)), self.range_do)
+        attp = tf.reshape(attp, (-1, 2, self.input_size, self.input_size))
         att  = atti + attp
         att  = tf.nn.softmax((att / math.sqrt(self.embed_dim)) + self.mask, axis=-1)
+        att1 = att[:, 0]
+        att2 = att[:, 1]
+        
+        # Differential attention
+        lambda_1 = tf.math.exp(tf.reduce_sum(self.lambda_q1 * self.lambda_k1, axis=-1))
+        lambda_2 = tf.math.exp(tf.reduce_sum(self.lambda_q2 * self.lambda_k2, axis=-1))
+        lambda_full = lambda_1 - lambda_2 + self.lambda_init
+        att = att1 - lambda_full * att2
+
         outi = att @ v
         attp = self.roll_embeddings(att, self.range_undo)
-        outp = attp @ pos
+        outp = attp @ pos_src
         out  = outi + outp
+        out = out * (1 - self.lambda_init)
         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)
+model = keras.models.load_model(
+    "rpc_diff_12b_320inp_ct4_01w10.keras",
+    custom_objects={
+        "DiffAttention" : DiffAttention,
+        "SharedEmbedding" : SharedEmbedding,
+        "masked_accuracy" : masked_accuracy
+    }
+)
+encoder = keras.Model(inputs=model.layers[0].input, outputs=model.layers[-1].output)
 encoder.summary()
 
 
@@ -166,10 +170,10 @@ 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")
+    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())
 
@@ -184,14 +188,15 @@ def generate(text, use_rpc=True, max_tokens=128):
         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]
+            _id = index.search(xq, size=1, epsilon=2)[0][0]
+            #_id = index.search(xq.reshape((1, -1)), 1)[1][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
+                if tmp in enc_text:
+                    tok = tmp
                 else: tok = all_toks[_id]
-            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))
@@ -199,6 +204,8 @@ def generate(text, use_rpc=True, max_tokens=128):
             tok = tf.argmax(res, axis=-1).numpy().tolist()
         
         enc_text += [tok]
-        response = tokenizer.decode(enc_text)
+        new_text = tokenizer.decode(enc_text)
+        res = new_text[len(text):]
+        text = new_text
     
-        yield response
+        yield res