Update app.py

app.py

@@ -5,21 +5,26 @@ import sys
 import urllib
 import json
 import torch
+from transformers import GPT2Tokenizer, GPT2LMHeadModel, GPT2Config
 
 def generate(tokenizer, model, text, features):
-    generated = tokenizer("<|startoftext|> <|titlestart|>{}<|titleend|>".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'],
     )
     for i, sample_output in enumerate(sample_outputs):
-        decoded = tokenizer.decode(sample_output, skip_special_tokens=True).replace('\\\\', '\\').split(text)[1]
-        st.write(decoded)
+        decoded = tokenizer.decode(sample_output, skip_special_tokens=False)
+        autor, text = decoded.split('<|authornamebegin|>')[1].split('<|authornameend|>')
+        st.markdown('**' + author.strip() + '**: ' + text.replace('<|endoftext|>', '').replace('<|pad|>', '').strip())
 
 
 def load_model():
     tokenizer = torch.load('./tokenizer.pt')
-    model = torch.load('./model.pt', map_location=torch.device('cpu'))
+    config = GPT2Config.from_json_file('./config.json')
+    model = GPT2LMHeadModel(config)
+    state_dict = torch.load('./pytorch_model.bin', map_location=torch.device('cpu'))
+    model.load_state_dict(state_dict)
     return tokenizer, model