Refactor to allow for non-ASCII data

generate.py

@@ -8,9 +8,9 @@
 from helpers import *
 from model import *
 
-def generate(decoder, prime_str='A', predict_len=100, temperature=0.8, cuda=False):
+def generate(decoder, vocab, prime_str='A', predict_len=100, temperature=0.8, cuda=False):
     hidden = decoder.init_hidden(1)
-    prime_input = Variable(char_tensor(prime_str).unsqueeze(0))
+    prime_input = Variable(char_tensor(prime_str, vocab).unsqueeze(0))
 
     if cuda:
         hidden = hidden.cuda()
@@ -31,9 +31,9 @@ def generate(decoder, prime_str='A', predict_len=100, temperature=0.8, cuda=Fals
         top_i = torch.multinomial(output_dist, 1)[0]
 
         # Add predicted character to string and use as next input
-        predicted_char = all_characters[top_i]
+        predicted_char = vocab[top_i]
         predicted += predicted_char
-        inp = Variable(char_tensor(predicted_char).unsqueeze(0))
+        inp = Variable(char_tensor(predicted_char, vocab).unsqueeze(0))
         if cuda:
             inp = inp.cuda()
 
@@ -51,7 +51,11 @@ def generate(decoder, prime_str='A', predict_len=100, temperature=0.8, cuda=Fals
     argparser.add_argument('--cuda', action='store_true')
     args = argparser.parse_args()
 
-    decoder = torch.load(args.filename)
+    try:
+        vocab, decoder = torch.load(args.filename)
+    except TypeError:
+        vocab = string.printable
+        decoder = torch.load(args.filename)
     del args.filename
-    print(generate(decoder, **vars(args)))
+    print(generate(decoder, vocab, **vars(args)))
 

helpers.py

@@ -1,28 +1,25 @@
 # https://github.com/spro/char-rnn.pytorch
 
-import unidecode
 import string
 import random
 import time
 import math
 import torch
 
-# Reading and un-unicode-encoding data
-
-all_characters = string.printable
-n_characters = len(all_characters)
+# Reading and extracting vocab from data
 
 def read_file(filename):
-    file = unidecode.unidecode(open(filename).read())
-    return file, len(file)
+    file = open(filename).read()
+    vocab = ''.join(sorted(set(file)))
+    return file, len(file), vocab
 
 # Turning a string into a tensor
 
-def char_tensor(string):
+def char_tensor(string, vocab):
     tensor = torch.zeros(len(string)).long()
     for c in range(len(string)):
         try:
-            tensor[c] = all_characters.index(string[c])
+            tensor[c] = vocab.index(string[c])
         except:
             continue
     return tensor

train.py

@@ -31,7 +31,7 @@
 if args.cuda:
     print("Using CUDA")
 
-file, file_len = read_file(args.filename)
+file, file_len, vocab = read_file(args.filename)
 
 def random_training_set(chunk_len, batch_size):
     inp = torch.LongTensor(batch_size, chunk_len)
@@ -40,8 +40,8 @@ def random_training_set(chunk_len, batch_size):
         start_index = random.randint(0, file_len - chunk_len)
         end_index = start_index + chunk_len + 1
         chunk = file[start_index:end_index]
-        inp[bi] = char_tensor(chunk[:-1])
-        target[bi] = char_tensor(chunk[1:])
+        inp[bi] = char_tensor(chunk[:-1], vocab)
+        target[bi] = char_tensor(chunk[1:], vocab)
     inp = Variable(inp)
     target = Variable(target)
     if args.cuda:
@@ -67,15 +67,15 @@ def train(inp, target):
 
 def save():
     save_filename = os.path.splitext(os.path.basename(args.filename))[0] + '.pt'
-    torch.save(decoder, save_filename)
+    torch.save((vocab, decoder), save_filename)
     print('Saved as %s' % save_filename)
 
 # Initialize models and start training
 
 decoder = CharRNN(
-    n_characters,
+    len(vocab),
     args.hidden_size,
-    n_characters,
+    len(vocab),
     model=args.model,
     n_layers=args.n_layers,
 )
@@ -97,7 +97,7 @@ def save():
 
         if epoch % args.print_every == 0:
             print('[%s (%d %d%%) %.4f]' % (time_since(start), epoch, epoch / args.n_epochs * 100, loss))
-            print(generate(decoder, 'Wh', 100, cuda=args.cuda), '\n')
+            print(generate(decoder, vocab, 'Wh', 100, cuda=args.cuda), '\n')
 
     print("Saving...")
     save()