Update app.py

app.py

@@ -17,10 +17,16 @@ def generate(tokenizer, model, text, features):
         )
         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|>')
+        if '<|authornamebegin|>' not in decoded:  
+            continue
+    
+        raw = decoded.split('<|authornamebegin|>')[-1]
+       
+        if '<|authornameend|>' not in raw:
+            continue
+      
+        author, text = raw.split('<|authornameend|>')
      
         count += 1
         st.markdown('**' + author.strip() + '**: ' + text.replace('<|endoftext|>', '').replace('<|pad|>', '').strip())