Added multi-vector embeddings

1 files changed, 13 insertions, 25 deletions

diff --git a/common.py b/common.py
index f369475..e8d3ac1 100644
--- a/common.py
+++ b/common.py
@@ -1,9 +1,10 @@
 from pathlib import Path
 import json
 
-import torch
+from models.clip.embeddings import ManagedCLIPTextEmbeddings
+from models.clip.tokenizer import MultiCLIPTokenizer
 
-from transformers import CLIPTextModel, CLIPTokenizer
+from safetensors import safe_open
 
 
 def load_config(filename):
@@ -18,33 +19,20 @@ def load_config(filename):
     return args
 
 
-def load_text_embedding(embeddings, token_id, file):
-    data = torch.load(file, map_location="cpu")
-
-    assert len(data.keys()) == 1, 'embedding data has multiple terms in it'
-
-    emb = next(iter(data.values()))
-    if len(emb.shape) == 1:
-        emb = emb.unsqueeze(0)
-
-    embeddings[token_id] = emb
-
-
-def load_text_embeddings(tokenizer: CLIPTokenizer, text_encoder: CLIPTextModel, embeddings_dir: Path):
+def load_embeddings_from_dir(tokenizer: MultiCLIPTokenizer, embeddings: ManagedCLIPTextEmbeddings, embeddings_dir: Path):
     if not embeddings_dir.exists() or not embeddings_dir.is_dir():
         return []
 
-    files = [file for file in embeddings_dir.iterdir() if file.is_file()]
-
-    tokens = [file.stem for file in files]
-    added = tokenizer.add_tokens(tokens)
-    token_ids = tokenizer.convert_tokens_to_ids(tokens)
-
-    text_encoder.resize_token_embeddings(len(tokenizer))
+    filenames = [filename for filename in embeddings_dir.iterdir() if filename.is_file()]
+    tokens = [filename.stem for filename in filenames]
 
-    token_embeds = text_encoder.get_input_embeddings().weight.data
+    for filename in embeddings_dir.iterdir():
+        if filename.is_file():
+            with safe_open(filename, framework="pt", device="cpu") as file:
+                embed = file.get_tensor("embed")
 
-    for (token_id, file) in zip(token_ids, files):
-        load_text_embedding(token_embeds, token_id, file)
+                added = tokenizer.add_multi_tokens(filename.stem, embed.shape[0])
+                embeddings.add_embed(added.placeholder_id)
+                embeddings.add_embed(added.multi_ids, embed)
 
     return tokens