Spaces:
Sleeping
Sleeping
File size: 1,350 Bytes
48ac659 |
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 |
from lib import *
from twokenize import tokenizeRawTweetText
import re
def muscaps_tokenize(raw):
raw = raw.lower()
for punc in string.punctuation:
raw = raw.replace(punc, ' ')
tokens = raw.split()
return tokens
def get_device(device_id: int) -> torch.device:
if not torch.cuda.is_available():
return torch.device('cpu')
device_id = min(torch.cuda.device_count() - 1, device_id)
return torch.device(f'cuda:{device_id}')
def preproc(caption, tokenizer, stop=True):
caption = caption.replace('.', '<STOP>')
caption_proc = tokenizer.encode(caption)
if stop:
caption_proc += tokenizer.encode('.')
return caption_proc
def postproc(caption):
caption = caption.replace('<STOP>', '.')
if caption[-1] == '.':
caption = caption[:-1]
return caption
class CheckpointManager:
def __init__(self):
self.checkpoint_dir = '/home/nsrivats/Repositories/MusicCaptioning/checkpoints'
def get_checkpoint(self, checkpoint):
with open(checkpoint, 'rb') as infile:
return torch.load(infile)
def save_checkpoint(self, state_dict, checkpoint):
filename = f'{self.checkpoint_dir}/{checkpoint}'
with open(filename, 'wb') as outfile:
torch.save(state_dict, outfile)
def save_logs(self, logdir):
pass
|