grantpitt commited on
Commit
5158d0e
1 Parent(s): 96f94cc

trying out pulling in CLIP

Browse files
Files changed (2) hide show
  1. pipeline.py +7 -4
  2. requirements.txt +5 -1
pipeline.py CHANGED
@@ -1,13 +1,15 @@
1
  from typing import Dict, List, Any
2
  import numpy as np
 
 
3
 
4
  class PreTrainedPipeline():
5
  def __init__(self, path=""):
6
- # IMPLEMENT_THIS
7
  # Preload all the elements you are going to need at inference.
8
  # For instance your model, processors, tokenizer that might be needed.
9
  # This function is only called once, so do all the heavy processing I/O here"""
10
- self.x = np.random.random(10)
 
11
 
12
  def __call__(self, inputs: str) -> List[float]:
13
  """
@@ -17,5 +19,6 @@ class PreTrainedPipeline():
17
  Return:
18
  A :obj:`list` of floats: The features computed by the model.
19
  """
20
- # IMPLEMENT_THIS
21
- return self.x.tolist()
 
 
1
  from typing import Dict, List, Any
2
  import numpy as np
3
+ from transformers import CLIPTokenizer, CLIPModel
4
+
5
 
6
  class PreTrainedPipeline():
7
  def __init__(self, path=""):
 
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.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
12
+ self.tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
13
 
14
  def __call__(self, inputs: str) -> List[float]:
15
  """
 
19
  Return:
20
  A :obj:`list` of floats: The features computed by the model.
21
  """
22
+ token_inputs = self.tokenizer([inputs], padding=True, return_tensors="pt")
23
+ query_embed = self.model.get_text_features(**token_inputs)
24
+ return query_embed.detach().cpu().numpy()[0].tolist()
requirements.txt CHANGED
@@ -1 +1,5 @@
1
- numpy==1.23.1
 
 
 
 
 
1
+ numpy==1.23.1
2
+ transformers==4.21.1
3
+ torch==1.12.1
4
+ torchvision==0.13.1
5
+ -f https://download.pytorch.org/whl/cu116