Option to edit records (#60)

* Rembed endpoint

* fixes token level embedding calculation

---------

Co-authored-by: FelixKirschKern <[email protected]>

Author: JWittmeyer

app.py

@@ -139,6 +139,16 @@ def upload_tensor_data(
     return responses.PlainTextResponse(status_code=status.HTTP_200_OK)
 
 
+@app.post("/re_embed_records/{project_id}")
+def re_embed_record(
+    project_id: str, request: data_type.EmbeddingRebuildRequest
+) -> responses.PlainTextResponse:
+    session_token = general.get_ctx_token()
+    controller.re_embed_records(project_id, request.changes)
+    general.remove_and_refresh_session(session_token)
+    return responses.PlainTextResponse(status_code=status.HTTP_200_OK)
+
+
 @app.put("/config_changed")
 def config_changed() -> responses.PlainTextResponse:
     config_handler.refresh_config()

controller.py

@@ -22,7 +22,7 @@
 from spacy.vocab import Vocab
 from data import data_type, doc_ock
 from embedders import Transformer
-from typing import Any, Dict, Iterator, List, Optional
+from typing import Any, Dict, Iterator, List, Optional, Union
 
 from util import daemon, request_util
 from util.config_handler import get_config_value
@@ -33,6 +33,8 @@
 import pandas as pd
 from submodules.s3 import controller as s3
 import openai
+import gc
+
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -77,7 +79,7 @@ def get_docbins(
         docs = list(doc_bin_loaded.get_docs(vocab))
         for col, doc in zip(record_item.columns, docs):
             if col == attribute_name:
-                result[record_item.record_id] = doc
+                result[str(record_item.record_id)] = doc
     result_list = []
     for record_id in record_ids_batch:
         result_list.append(result[record_id])
@@ -409,14 +411,14 @@ def run_encoding(
         if embedding_type == enums.EmbeddingType.ON_ATTRIBUTE.value:
             request_util.post_embedding_to_neural_search(project_id, embedding_id)
 
-        if get_config_value("is_managed"):
-            pickle_path = os.path.join(
-                "/inference", project_id, f"embedder-{embedding_id}.pkl"
-            )
-            if not os.path.exists(pickle_path):
-                os.makedirs(os.path.dirname(pickle_path), exist_ok=True)
-                with open(pickle_path, "wb") as f:
-                    pickle.dump(embedder, f)
+        # now always since otherwise record edit wouldn't work for embedded columns
+        pickle_path = os.path.join(
+            "/inference", project_id, f"embedder-{embedding_id}.pkl"
+        )
+        if not os.path.exists(pickle_path):
+            os.makedirs(os.path.dirname(pickle_path), exist_ok=True)
+            with open(pickle_path, "wb") as f:
+                pickle.dump(embedder, f)
 
         upload_embedding_as_file(project_id, embedding_id)
         embedding.update_embedding_state_finished(
@@ -499,3 +501,96 @@ def upload_embedding_as_file(
 
 def __is_embedders_internal_model(model_name: str):
     return model_name in ["bag-of-characters", "bag-of-words", "tf-idf"]
+
+
+def re_embed_records(project_id: str, changes: Dict[str, List[Dict[str, str]]]):
+    for embedding_id in changes:
+        if len(changes[embedding_id]) == 0:
+            continue
+
+        embedding_item = embedding.get(project_id, embedding_id)
+        if not embedding_item:
+            continue
+
+        # convert to int since the request automatically converts it to string
+        if "sub_key" in changes[embedding_id][0]:
+            for d in changes[embedding_id]:
+                d["sub_key"] = int(d["sub_key"])
+
+        embedder = __setup_tmp_embedder(project_id, embedding_id)
+
+        data_to_embed = None
+        record_ids = None  # Either list or set depending on embedding type
+        attribute_name = changes[embedding_id][0]["attribute_name"]
+
+        if embedding_item.type == enums.EmbeddingType.ON_TOKEN.value:
+            # can't have sub_key so records are unique so we can just get them all since order is preserved in get_docbins
+            record_ids = [c["record_id"] for c in changes[embedding_id]]
+            data_to_embed = get_docbins(
+                project_id, record_ids, embedder.nlp.vocab, attribute_name
+            )
+        else:
+            # order is important, data collection request doesn't order so we do it ourselves
+            record_ids = {c["record_id"] for c in changes[embedding_id]}
+            records = record.get_by_record_ids(project_id, record_ids)
+            records = {str(r.id): r for r in records}
+
+            data_to_embed = [
+                records[c["record_id"]].data[attribute_name]
+                if "sub_key" not in c
+                else records[c["record_id"]].data[attribute_name][c["sub_key"]]
+                for c in changes[embedding_id]
+            ]
+
+        new_tensors = embedder.transform(data_to_embed)
+
+        if len(new_tensors) != len(changes[embedding_id]):
+            raise Exception(
+                f"Number of new tensors ({len(new_tensors)}) doesn't match number of changes ({len(changes[embedding_id])})"
+            )
+
+        # delete old
+        if "sub_key" in changes[embedding_id][0]:
+            embedding.delete_by_record_ids_and_sub_keys(
+                project_id,
+                embedding_id,
+                [(c["record_id"], c["sub_key"]) for c in changes[embedding_id]],
+            )
+        else:
+            embedding.delete_by_record_ids(project_id, embedding_id, record_ids)
+        # add new
+        record_ids_batched = [
+            c["record_id"]
+            if "sub_key" not in c
+            else c["record_id"] + "@" + str(c["sub_key"])
+            for c in changes[embedding_id]
+        ]
+
+        embedding.create_tensors(
+            project_id,
+            embedding_id,
+            record_ids_batched,
+            new_tensors,
+            with_commit=True,
+        )
+
+        upload_embedding_as_file(project_id, embedding_id)
+        request_util.delete_embedding_from_neural_search(embedding_id)
+        request_util.post_embedding_to_neural_search(project_id, embedding_id)
+
+        del embedder
+        time.sleep(0.1)
+        gc.collect()
+        time.sleep(0.1)
+
+
+def __setup_tmp_embedder(project_id: str, embedder_id: str) -> Transformer:
+    embedder_path = os.path.join(
+        "/inference", project_id, f"embedder-{embedder_id}.pkl"
+    )
+    if not os.path.exists(embedder_path):
+        raise Exception(f"Embedder {embedder_id} not found")
+    with open(embedder_path, "rb") as f:
+        embedder = pickle.load(f)
+
+    return embedder

data/data_type.py

@@ -1,9 +1,16 @@
-from dataclasses import dataclass
-from typing import Optional
+from typing import Dict, List
 from pydantic import BaseModel
 
 
 class EmbeddingRequest(BaseModel):
     project_id: str
     embedding_id: str
 
+
+class EmbeddingRebuildRequest(BaseModel):
+    # example request structure:
+    # {"<embedding_id>":[{"record_id":"<record_id>","attribute_name":"<attribute_name>","sub_key":<sub_key>}]}
+    # note that sub_key is optional and only for embedding lists relevant
+    # also sub_key is an int but converted to string in the request
+
+    changes: Dict[str, List[Dict[str, str]]]