-
Notifications
You must be signed in to change notification settings - Fork 5
/
utils.py
28 lines (22 loc) · 759 Bytes
/
utils.py
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
#!/usr/bin/env python
"""Defines some useful functions
"""
import os
import numpy as np
import nltk
__author__ = "jssprz"
__version__ = "0.0.1"
__maintainer__ = "jssprz"
__email__ = "[email protected]"
__status__ = "Development"
def get_freer_gpu():
os.system('nvidia-smi -q -d Memory |grep -A4 GPU|grep Free >tmp')
memory_available = [int(x.split()[2]) for x in open('tmp', 'r').readlines()]
return np.argsort(memory_available)[::-1]
def get_tags_feat(captions, tags_vocab):
result = np.zeros((len(tags_vocab),), dtype='int')
for caption in captions:
tags = [x for x in nltk.tokenize.word_tokenize(caption.lower()) if x in key_words]
for tag in tags:
result[tags_vocab(tag)] = 1
return result