Commit
•
b7332d3
1
Parent(s):
ec7cd57
Added inference pipeline
Browse files- pipeline.py +31 -0
pipeline.py
CHANGED
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
|
2 |
+
import io
|
3 |
+
import base64
|
4 |
+
import shutil
|
5 |
+
from transformers import CLIPProcessor, CLIPModel, CLIPTokenizer
|
6 |
+
|
7 |
+
|
8 |
+
|
9 |
+
class PreTrainedPipeline():
|
10 |
+
def __init__(self, path=""):
|
11 |
+
"""
|
12 |
+
Initialize model
|
13 |
+
"""
|
14 |
+
self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
|
15 |
+
self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
|
16 |
+
self.tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
|
17 |
+
|
18 |
+
def __call__(self, inputs: str):
|
19 |
+
"""
|
20 |
+
Args:
|
21 |
+
inputs (:obj:`str`):
|
22 |
+
a string containing some text
|
23 |
+
Return:
|
24 |
+
A :obj:`list`list of floats: The features computed by the model.
|
25 |
+
"""
|
26 |
+
inputs = self.tokenizer(inputs, padding=True, return_tensors="pt")
|
27 |
+
|
28 |
+
# Compute text embeddings
|
29 |
+
with torch.no_grad():
|
30 |
+
text_features = self.model.get_text_features(**inputs)
|
31 |
+
return text_features[0].tolist()
|