feat: add additional params for reducing repetition

app.py

@@ -29,8 +29,7 @@ class DefaultArgs:
 
 if os.getenv("RUNNING_ON_HF_SPACE"):
     login(token=os.getenv("HF_TOKEN"))
-    hf_repo = "leemeng/stablelm-jp-alpha"
-    
+    hf_repo = os.getenv("HF_MODEL_REPO")
     args = DefaultArgs()
     args.hf_model_name_or_path = hf_repo
     args.spm_model_path = hf_hub_download(repo_id=hf_repo, filename="sentencepiece.model")
@@ -86,35 +85,35 @@ class SentencePieceStreamer(BaseStreamer):
             return
         
         self.generated_text += text
-        # print(f"[streamer]: {self.generated_text}")
-        # yield text
+        logging.debug(f"[streamer]: {self.generated_text}")
     
     def end(self):
         self.ended = True
 
-def user(user_message, history):
-    logging.debug(f"[user] user_message: {user_message}")
+def user(prompt, history):
+    logging.info(f"[user] prompt: {prompt}")
     logging.debug(f"[user] history: {history}")
     
-    res = ("", history + [[user_message, None]])
+    res = ("", history + [[prompt, None]])
     return res
 
 def bot(
     history, 
+    do_sample,
     temperature,
+    repetition_penalty,
+    no_repeat_ngram_size,
     max_new_tokens,
 ):
-    logging.debug(f"[bot] history: {history}")
-    logging.debug(f"temperature: {temperature}")
+    logging.info("[bot]")
+    logging.info(dict(locals()))
+    logging.debug(f"history: {history}")
     
     # TODO: modify `<br>` back to `\n` based on the original user prinpt
     prompt = history[-1][0]
     
     tokens = sp.encode(prompt)
     input_ids = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(model.device)
-    
-    # TODO: parametrize setting on UI
-    do_sample = True
         
     streamer = SentencePieceStreamer(sp=sp)
     
@@ -124,14 +123,15 @@ def bot(
     
     thr = threading.Thread(target=model.generate, args=(), kwargs=dict(
         input_ids=input_ids,
+        do_sample=do_sample,
         temperature=temperature,
+        repetition_penalty=repetition_penalty,
+        no_repeat_ngram_size=no_repeat_ngram_size,
         max_new_tokens=max_possilbe_new_tokens,
-        do_sample=do_sample,
         streamer=streamer,
         # max_length=4096,
         # top_k=100,
         # top_p=0.9,
-        # repetition_penalty=1.0,
         # num_return_sequences=2,
         # num_beams=2,
     ))
@@ -145,6 +145,7 @@ def bot(
     
     # TODO: optimize for final few tokens
     history[-1][1] = streamer.generated_text
+    logging.info(f"generation: {history[-1][1]}")
     yield history
 
 if gr_interface:
@@ -154,8 +155,13 @@ with gr.Blocks() as gr_interface:
     chatbot = gr.Chatbot(label="StableLM JP Alpha").style(height=500)
     
     # generation params
+    do_sample = gr.Checkbox(True, label="Do Sample", visible=False)
+    
     with gr.Row():
         temperature = gr.Slider(0, 1, value=0.7, step=0.05, label="Temperature")
+        repetition_penalty = gr.Slider(1, 1.5, value=1.2, step=0.05, label="Repetition Penalty")
+    with gr.Row():
+        no_repeat_ngram_size = gr.Slider(0, 10, value=5, step=1, label="No Repeat Ngram Size")
         max_new_tokens = gr.Slider(
             128, 
             model.config.max_position_embeddings, 
@@ -191,18 +197,27 @@ with gr.Blocks() as gr_interface:
     with gr.Row():
         submit = gr.Button("Submit")
         stop = gr.Button("Stop")
-    
-    clear = gr.Button("Clear History")
+        clear = gr.Button("Clear")
     
     # event handling
-    submit_event = prompt.submit(user, [prompt, chatbot], [prompt, chatbot], queue=False)\
-        .then(bot, [chatbot, temperature, max_new_tokens], chatbot, queue=True)
-    
-    submit_click_event = submit.click(user, [prompt, chatbot], [prompt, chatbot], queue=False)\
-        .then(bot, [chatbot, temperature, max_new_tokens], chatbot, queue=True)
+    user_io = [prompt, chatbot]
+    bot_inputs = [
+        chatbot, 
+        do_sample, 
+        temperature, 
+        repetition_penalty, 
+        no_repeat_ngram_size,
+        max_new_tokens,
+    ]
+    
+    submit_event = prompt.submit(user, user_io, user_io, queue=False)\
+        .then(bot, bot_inputs, chatbot, queue=True)
+    
+    submit_click_event = submit.click(user, user_io, user_io, queue=False)\
+        .then(bot, bot_inputs, chatbot, queue=True)
     
     stop.click(None, None, None, cancels=[submit_event, submit_click_event], queue=False)
     clear.click(lambda: None, None, chatbot, queue=False)
 
-gr_interface.queue()
+gr_interface.queue(max_size=32, concurrency_count=2)
 gr_interface.launch(server_port=args.port, share=args.make_public)