import os
import torch
from torch.autograd import Variable
from torch import optim
import torch.nn.functional as F
import torch.nn as nn
import argparse
import models
import math
parser = argparse.ArgumentParser(description='sample.py')
parser.add_argument('-init', default='The meaning of life is ',
help="""Initial text """)
parser.add_argument('-load_model', default='',
help="""Model filename to load""")
parser.add_argument('-seq_length', type=int, default=50,
help="""Maximum sequence length""")
parser.add_argument('-temperature', type=float, default=0.4,
help="""Temperature for sampling.""")
parser.add_argument('-neuron', type=int, default=0,
help="""Neuron to read.""")
parser.add_argument('-overwrite', type=float, default=0,
help="""Value used to overwrite the neuron. 0 means don't overwrite.""")
parser.add_argument('-layer', type=int, default=-1,
help="""Layer to read. -1 = last layer""")
# GPU
parser.add_argument('-cuda', action='store_true',
help="""Use CUDA""")
opt = parser.parse_args()
def batchify(data, bsz):
tokens = len(data.encode())
ids = torch.LongTensor(tokens)
token = 0
for char in data.encode():
ids[token] = char
token += 1
nbatch = ids.size(0) // bsz
ids = ids.narrow(0, 0, nbatch * bsz)
ids = ids.view(bsz, -1).t().contiguous()
return ids
def color(p):
p = math.tanh(3*p)*.5+.5
q = 1.-p*1.3
r = 1.-abs(0.5-p)*1.3+.3*q
p=1.3*p-.3
i = int(p*255)
j = int(q*255)
k = int(r*255)
if j<0:
j=0
if k<0:
k=0
if k >255:
k=255
if i<0:
i = 0
return ('\033[38;2;%d;%d;%dm' % (j, k, i)).encode()
batch_size = 1
checkpoint = torch.load(opt.load_model)
embed = checkpoint['embed']
rnn = checkpoint['rnn']
loss_fn = nn.CrossEntropyLoss()
text = batchify(opt.init, batch_size)
def make_cuda(state):
if isinstance(state, tuple):
return (state[0].cuda(), state[1].cuda())
else:
return state.cuda()
batch = Variable(text)
states = rnn.state0(batch_size)
if isinstance(states, tuple):
hidden, cell = states
else:
hidden = states
last = hidden.size(0)-1
if opt.layer <= last and opt.layer >= 0:
last = opt.layer
if opt.cuda:
batch =batch.cuda()
states = make_cuda(states)
embed.cuda()
rnn.cuda()
loss_avg = 0
loss = 0
gen = bytearray()
for t in range(text.size(0)):
emb = embed(batch[t])
ni = (batch[t]).data[0]
states, output = rnn(emb, states)
if isinstance(states, tuple):
hidden, cell = states
else:
hidden = states
feat = hidden.data[last,0,opt.neuron]
if ni< 128:
col = color(feat)
gen+=(col)
gen.append(ni)
print(opt.init)
if opt.temperature == 0:
topv, topi = output.data.topk(1)
ni = topi[0][0]
gen.append(ni)
inp = Variable(topi[0], volatile=True)
if opt.cuda:
inp = inp.cuda()
for t in range(opt.seq_length):
emb = embed(inp)
states, output = rnn(emb, states)
topv, topi = output.data.topk(1)
ni = topi[0][0]
gen.append(ni)
inp = Variable(topi[0])
if opt.cuda:
inp = inp.cuda()
else:
probs = F.softmax(output[0].squeeze().div(opt.temperature)).data.cpu()
ni = torch.multinomial(probs,1)[0]
feat = hidden.data[last,0,opt.neuron]
if ni < 128:
col = color(feat)
gen+=(col)
gen.append(ni)
inp = Variable(torch.LongTensor([ni]), volatile=True)
if opt.cuda:
inp = inp.cuda()
for t in range(opt.seq_length):
emb = embed(inp)
states, output = rnn(emb, states)
if isinstance(states, tuple):
hidden, cell = states
else:
hidden = states
feat = hidden.data[last,0,opt.neuron]
if isinstance(output, list):
output =output[0]
probs = F.softmax(output.squeeze().div(opt.temperature)).data.cpu()
ni = torch.multinomial(probs,1)[0]
if ni< 128:
col = color(feat)
gen+=(col)
gen.append(ni)
inp = Variable(torch.LongTensor([ni]))
if opt.cuda:
inp = inp.cuda()
if opt.overwrite != 0:
hidden.data[last,0,opt.neuron] = opt.overwrite
gen+=('\033[0m').encode()
print(gen.decode("utf-8",errors = 'ignore' ))
