Update rpc.py

rpc.py

@@ -149,16 +149,13 @@ class DiffAttention(keras.layers.Layer):
         att = att1 - lambda_full * att2
 
         outi = att @ v
-        attp = self.roll_embeddings(att, self.range_undo)
-        outp = attp @ pos_src
-        out  = outi + outp
         out = out * (1 - self.lambda_init)
         return out
     
 
 # Import Model
 model = keras.models.load_model(
-    "rpc_diff_12b_320inp_ct4_01w10.keras",
+    "rpc.keras",
     custom_objects={
         "DiffAttention" : DiffAttention,
         "SharedEmbedding" : SharedEmbedding,
@@ -187,13 +184,20 @@ def vectorize_texts(all_texts):
 # Import Database and All Toks
 index = None
 all_toks = None
-def load_index(index_path="/dev/shm/rpc-vecdb/index"):
+index_type = None
+def load_index(index_path="/dev/shm/rpc-vecdb/index", idx_type="ngt"):
     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")
+    global index_type
+    index_type = idx_type
+    if idx_type == "ngt":
+        import ngtpy
+        index = ngtpy.Index(index_path, read_only=True)
+    elif idx_type == "faiss":    
+        import faiss
+        index = faiss.read_index(index_path + "/index.faiss")
+    else:
+        raise ValueError("Unknown index type")
     with open(index_path + "/all_toks.json", "r") as f:
         all_toks = json.loads(f.read())
 
@@ -209,8 +213,10 @@ 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=1)[0][0]
-            #_id = index.search(xq.reshape((1, -1)), 1)[1][0][0]
+            if index_type == "ngt":
+                _id = index.search(xq, size=1, epsilon=1)[0][0]
+            else:
+                _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 tmp in enc_text: