bryandts commited on
Commit
423fc3f
1 Parent(s): 98688ba

Update preprocessing.py

Browse files
Files changed (1) hide show
  1. preprocessing.py +9 -0
preprocessing.py CHANGED
@@ -8,6 +8,15 @@ import torch
8
  from torch.autograd import Variable
9
  import pdb
10
  import torch.nn.functional as F
 
 
 
 
 
 
 
 
 
11
 
12
  class Text2ImageDataset(Dataset):
13
 
 
8
  from torch.autograd import Variable
9
  import pdb
10
  import torch.nn.functional as F
11
+ from sentence_transformers import SentenceTransformer
12
+
13
+ model = SentenceTransformer('sentence-transformers/all-MiniLM-L12-v2')
14
+
15
+ def process_caption(text):
16
+ # Encode text and get model output
17
+ embeddings = model.encode(text)
18
+ ouput = np.expand_dims(embeddings, axis=0)
19
+ return ouput
20
 
21
  class Text2ImageDataset(Dataset):
22