bug fixes in update_text

rishiraj · Apr 14, 2024 · 812af60 · 812af60
1 parent 6674a2b
commit 812af60
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/spanking/main.py b/spanking/main.py
@@ -26,8 +26,8 @@ def delete_text(self, index):
     def update_text(self, index, new_text):
         if 0 <= index < len(self.texts):
             self.texts[index] = new_text
-            new_embedding = self.model.encode([new_text], normalize_embeddings=True)
-            self.embeddings = jax.ops.index_update(self.embeddings, index, new_embedding)
+            new_embedding = self.model.encode([new_text], normalize_embeddings=True).squeeze()
+            self.embeddings = (self.embeddings).at[index].set(new_embedding)
         else:
             raise IndexError("Invalid index")