attept 1 of actual ranking
Browse files- pipeline.py +18 -4
pipeline.py
CHANGED
@@ -1,6 +1,7 @@
|
|
1 |
from typing import Dict, List, Any
|
2 |
import numpy as np
|
3 |
from transformers import CLIPTokenizer, CLIPModel
|
|
|
4 |
|
5 |
|
6 |
class PreTrainedPipeline():
|
@@ -8,10 +9,15 @@ class PreTrainedPipeline():
|
|
8 |
# Preload all the elements you are going to need at inference.
|
9 |
# For instance your model, processors, tokenizer that might be needed.
|
10 |
# This function is only called once, so do all the heavy processing I/O here"""
|
11 |
-
self.
|
12 |
-
self.
|
13 |
|
14 |
-
|
|
|
|
|
|
|
|
|
|
|
15 |
"""
|
16 |
Args:
|
17 |
inputs (:obj:`str`):
|
@@ -21,4 +27,12 @@ class PreTrainedPipeline():
|
|
21 |
"""
|
22 |
token_inputs = self.tokenizer([inputs], padding=True, return_tensors="pt")
|
23 |
query_embed = self.model.get_text_features(**token_inputs)
|
24 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
from typing import Dict, List, Any
|
2 |
import numpy as np
|
3 |
from transformers import CLIPTokenizer, CLIPModel
|
4 |
+
import os
|
5 |
|
6 |
|
7 |
class PreTrainedPipeline():
|
|
|
9 |
# Preload all the elements you are going to need at inference.
|
10 |
# For instance your model, processors, tokenizer that might be needed.
|
11 |
# This function is only called once, so do all the heavy processing I/O here"""
|
12 |
+
self.sign_ids = np.load(os.path.join(path, "sign_ids.npy"))
|
13 |
+
self.sign_embeddings = np.load(os.path.join(path, "vanilla_large-patch14_image_embeddings.npy"))
|
14 |
|
15 |
+
hf_model_path = "openai/clip-vit-large-patch14"
|
16 |
+
self.model = CLIPModel.from_pretrained(hf_model_path)
|
17 |
+
self.tokenizer = CLIPTokenizer.from_pretrained(hf_model_path)
|
18 |
+
|
19 |
+
|
20 |
+
def __call__(self, inputs: str):
|
21 |
"""
|
22 |
Args:
|
23 |
inputs (:obj:`str`):
|
|
|
27 |
"""
|
28 |
token_inputs = self.tokenizer([inputs], padding=True, return_tensors="pt")
|
29 |
query_embed = self.model.get_text_features(**token_inputs)
|
30 |
+
np_query_embed = query_embed.detach().cpu().numpy()[0]
|
31 |
+
|
32 |
+
# Compute the cosine similarity; note the embeddings are normalized.
|
33 |
+
cos_similarites = self.sign_embeddings @ np_query_embed
|
34 |
+
sign_id_arg_rankings = np.argsort(cos_similarites)[::-1]
|
35 |
+
n = 50
|
36 |
+
top_sign_ids = self.sign_ids[sign_id_arg_rankings[:n]]
|
37 |
+
top_sign_similarities = cos_similarites[sign_id_arg_rankings[:n]]
|
38 |
+
return dict(zip(top_sign_ids.tolist(), top_sign_similarities.tolist()))
|