Fix new tokenizer problem

RobinDong · May 2, 2024 · f724c26 · f724c26
1 parent 4174dab
commit f724c26
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/tinymm/CLIP/demo.py b/tinymm/CLIP/demo.py
@@ -56,7 +56,7 @@ def start(self, checkpoint: str, cuda: bool = False):
         text = st.text_input("Input:")
         if not text:
             return
-        ids = self.enc(text)
+        ids = self.enc(text)["input_ids"]
         ids = np.pad(ids, (0, (self.seq_len - len(ids))), "constant")
         ids = torch.tensor(ids).unsqueeze(0)
         txt_embd = model.txt_encoder(ids)

diff --git a/tinymm/CLIP/evaluate.py b/tinymm/CLIP/evaluate.py
@@ -39,7 +39,7 @@ def __init__(self, model):
         print("Compute embeddings for all categories...")
         for index in tqdm(range(1, self.nr_categories + 1)):
             text = f"A photo of a {self.cats[index]}"
-            ids = enc(text)
+            ids = enc(text)["input_ids"]
             ids = np.pad(ids, (0, (self.seq_len - len(ids))), "constant")
             ids = torch.tensor(ids).unsqueeze(0)
             with torch.no_grad():