Upload get_embeds.py

get_embeds.py

@@ -0,0 +1,32 @@
+from transformers import Owlv2TextModel, Owlv2Processor, AutoTokenizer
+import json
+import torch
+from torch import nn
+import tqdm
+
+embed_dict = nn.ParameterDict()
+bsz = 8
+
+with open("id_to_str.json") as f:
+    data = json.load(f)
+
+keys = list(data.keys())
+bar = tqdm.tqdm(range(len(keys)//bsz))
+
+proc = Owlv2Processor.from_pretrained("google/owlv2-base-patch16-ensemble")
+tokenizer = AutoTokenizer.from_pretrained("google/owlv2-base-patch16-ensemble")
+model = Owlv2TextModel.from_pretrained("google/owlv2-base-patch16-ensemble")
+
+for i in bar:
+    batch = [data[key].replace("_", " ") for key in keys[i*bsz:(i+1)*bsz]]
+    tokenized = tokenizer(batch)
+    for k in range(bsz):
+        if len(tokenized[k]) > 16:
+            tokenizer.decode(tokenized[k])
+
+    batch = proc(text=batch, return_tensors="pt")
+    output = model(**batch)
+    for k, key in enumerate(keys[i*bsz:(i+1)*bsz]):
+        embed_dict[key] = output.pooler_output[k, :]
+
+torch.save(embed_dict.state_dict(), "embeds.pt")