Update app.py

app.py

@@ -9,6 +9,20 @@ with open('config.yaml', "r") as f:
 tokenizer = AutoTokenizer.from_pretrained("Xibanya/DS9Bot")
 model = AutoModelForCausalLM.from_pretrained("Xibanya/DS9Bot")
 
+def split_dialogue(line):
+    tokens = line.split()
+    concat = ''
+    total = len(tokens)
+    for t in range(total):
+        token = tokens[t]
+        if token != '[OC]:' and '[on ' not in token and \
+                ('[' in token or
+                 (token.isupper() and (':' in token or (t < total - 1 and tokens[t + 1] == '[OC]:')))):
+            token = '\n' + token
+
+        concat = concat + token + ' '
+    return concat.strip()
+
 def generate(prompt: str = ' ', temp: float = 0.7, length: float = 250):
     torch.Generator().manual_seed(cfg['seed'] if cfg['seed'] is not None else torch.seed())
     prompt = '[Ops] SISKO: ' if prompt is None or '' else prompt
@@ -27,9 +41,7 @@ def generate(prompt: str = ' ', temp: float = 0.7, length: float = 250):
             )
     text = tokenizer.batch_decode(
             output, skip_special_tokens=True, clean_up_tokenization_spaces=True)[0]
-    for string in cfg['split']:
-            text = text.replace(string, '\n' + string.strip())
-    return text.strip()
+    return split_dialogue(text)
     
 title = 'Deep Space Nine Script Generator'