[bug]

.gitignore

@@ -4,3 +4,5 @@
 
 **/flagged/
 **/__pycache__/
+
+trained_models/

main.py

@@ -3,6 +3,7 @@
 import argparse
 from collections import defaultdict
 import os
+import platform
 
 import gradio as gr
 from threading import Thread
@@ -71,7 +72,6 @@ def main():
         input_ids = torch.tensor([input_ids], dtype=torch.long)
         input_ids = input_ids.to(device)
 
-        output: str = ""
         streamer = TextIteratorStreamer(tokenizer=tokenizer)
 
         generation_kwargs = dict(
@@ -88,17 +88,24 @@ def main():
         thread = Thread(target=model.generate, kwargs=generation_kwargs)
         thread.start()
 
+        output: str = ""
+        first_answer = True
         for output_ in streamer:
-            output_ = output_.replace(" ", "")
-            output_ = output_.replace("[CLS]", "")
+            if first_answer:
+                first_answer = False
+                continue
+            # output_ = output_.replace(text, "")
+            # output_ = output_.replace("[CLS]", "")
             output_ = output_.replace("[SEP]", "\n")
             output_ = output_.replace("[UNK]", "")
-            output_ = output_.replace(text, "")
+            output_ = output_.replace(" ", "")
 
             output += output_.strip()
             output_text_box.value += output
             yield output
 
+    model_name_choices = ["trained_models/lib_service_4chan"] \
+        if platform.system() == "Windows" else ["qgyd2021/lib_service_4chan"]
     demo = gr.Interface(
         fn=fn_stream,
         inputs=[
@@ -107,7 +114,7 @@ def main():
             gr.Slider(minimum=0, maximum=1, value=0.85, step=0.01, label="top_p"),
             gr.Slider(minimum=0, maximum=1, value=0.35, step=0.01, label="temperature"),
             gr.Slider(minimum=0, maximum=2, value=1.2, step=0.01, label="repetition_penalty"),
-            gr.Dropdown(choices=["qgyd2021/lib_service_4chan"], value="qgyd2021/lib_service_4chan", label="model_name"),
+            gr.Dropdown(choices=model_name_choices, value=model_name_choices[0], label="model_name"),
             gr.Checkbox(value=True, label="is_chat")
         ],
         outputs=[output_text_box],