Update app.py

app.py

@@ -8,7 +8,7 @@ import torch
 from transformers import GPT2Tokenizer, GPT2LMHeadModel, GPT2Config
 
 def generate(tokenizer, model, text, features):
-    generated = tokenizer("<|startoftext|> <|titlestart|>{}<|titleend|><|authornamebegin|>".format(text), return_tensors="pt").input_ids
+    generated = tokenizer("<|startoftext|><|titlestart|>{}<|titleend|><|authornamebegin|>".format(text), return_tensors="pt").input_ids
     sample_outputs = model.generate(
         generated, do_sample=True, top_k=50, 
         max_length=features['max_length'], top_p=features['top_p'], temperature=features['t'] / 100.0, num_return_sequences=features['num'],
@@ -20,7 +20,10 @@ def generate(tokenizer, model, text, features):
 
 
 def load_model():
-    tokenizer = torch.load('./tokenizer.pt')
+    additional_special_tokens = ['<|titlestart|>', '<|titleend|>', '<|authornamebegin|>', '<|authornameend|>']
+    tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium', bos_token='<|startoftext|>',
+                                          eos_token='<|endoftext|>', pad_token='<|pad|>',
+                                          additional_special_tokens=additional_special_tokens)
     config = GPT2Config.from_json_file('./config.json')
     model = GPT2LMHeadModel(config)
     state_dict = torch.load('./pytorch_model.bin', map_location=torch.device('cpu'))