Spaces:
Sleeping
Sleeping
Update preprocessing.py
Browse files- 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 |
|