Update app.py

app.py

@@ -25,7 +25,8 @@ def rnd_generate(prompt):
                        temperature=0.75,
                        top_k=50, 
                        top_p=0.9)
-    return tokenizer.decode(results[0])
+    result = tokenizer.decode(results[0]).split("<bos>")[1].split("<eos>")[0]
+    return result
 
 @spaces.GPU
 def generate(prompt):
@@ -43,7 +44,10 @@ def generate(prompt):
                        top_k=50, 
                        top_p=0.9
                       )
-    result = tokenizer.decode(results[0]).split("<bos>")[1]
+    try:
+        result = tokenizer.decode(results[0]).split("<bos>")[1].split("<eos>")[0]
+    except IndexError:
+        result = tokenizer.decode(results[0]).split("<bos>")[1]
     return result
 
 examples = [