File size: 1,503 Bytes
f9567e5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
"""
This file is used to extract feature of the empty prompt.
"""
import os
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import torch
import os
import numpy as np
from libs.clip import FrozenCLIPEmbedder
from libs.t5 import T5Embedder
def main():
prompts = [
'',
]
device = 'cuda'
llm = 'clip'
if llm=='clip':
clip = FrozenCLIPEmbedder()
clip.eval()
clip.to(device)
elif llm=='t5':
t5 = T5Embedder(device=device)
else:
raise NotImplementedError
save_dir = f'./'
if llm=='clip':
latent, latent_and_others = clip.encode(prompts)
token_embedding = latent_and_others['token_embedding']
token_mask = latent_and_others['token_mask']
token = latent_and_others['tokens']
elif llm=='t5':
latent, latent_and_others = t5.get_text_embeddings(prompts)
token_embedding = latent_and_others['token_embedding'].to(torch.float32) * 10.0
token_mask = latent_and_others['token_mask']
token = latent_and_others['tokens']
for i in range(len(prompts)):
data = {'token_embedding': token_embedding[i].detach().cpu().numpy(),
'token_mask': token_mask[i].detach().cpu().numpy(),
'token': token[i].detach().cpu().numpy(),
'batch_caption': prompts[i]}
np.save(os.path.join(save_dir, f'empty_context.npy'), data)
if __name__ == '__main__':
main()
|