Update app.py

app.py

@@ -9,15 +9,20 @@ 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
-    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=False)
-        if '<|authornameend|>' not in decoded:
+    count = 0
+    while count < features['num']:
+        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=1,
+        )
+        decoded = tokenizer.decode(sample_outputs[0], skip_special_tokens=False)
+        print(decoded, file=sys.stderr)
+        if '<|authornameend|>' not in decoded:  
             continue
-        author, text = decoded.split('<|authornamebegin|>')[1].split('<|authornameend|>')
+        
+        author, text = decoded.split('<|authornamebegin|>')[-1].split('<|authornameend|>')
+     
+        count += 1
         st.markdown('**' + author.strip() + '**: ' + text.replace('<|endoftext|>', '').replace('<|pad|>', '').strip())