Update app.py

app.py

@@ -33,9 +33,12 @@ class StarlingBot:
     def __init__(self, system_prompt="I am Starling-7B by Tonic-AI, I ready to do anything to help my user."):
         self.system_prompt = system_prompt
 
-    def predict(self, user_message, assistant_message, system_prompt, do_sample, temperature=0.4, max_new_tokens=700, top_p=0.99, repetition_penalty=1.9):
+    def predict(self, user_message, assistant_message, system_prompt, mode, do_sample, temperature=0.4, max_new_tokens=700, top_p=0.99, repetition_penalty=1.9):
         try:
-            conversation = f" GPT4 Correct Assistant: {system_prompt if system_prompt else self.system_prompt} <|end_of_turn|> GPT4 Correct Assistant: {assistant_message if assistant_message else ''} <|end_of_turn|> GPT4 Correct User: {user_message} <|end_of_turn|> GPT4 Correct Assistant:"
+            if mode == "Assistant":
+                conversation = f" GPT4 Correct Assistant: {system_prompt if system_prompt else self.system_prompt}  GPT4 Correct Assistant: {assistant_message if assistant_message else ''}  GPT4 Correct User: {user_message}  GPT4 Correct Assistant:"
+            else:  # mode == "Coder"
+                conversation = f" Code Assistant: {system_prompt if system_prompt else self.system_prompt}  Code Assistant: {assistant_message if assistant_message else ''}  GPT4 Correct User: {user_message}  Code Assistant:"
             input_ids = tokenizer.encode(conversation, return_tensors="pt", add_special_tokens=True)
             input_ids = input_ids.to(device)
             response = model.generate(
@@ -84,8 +87,8 @@ with gr.Blocks(theme="ParityError/Anime") as demo:
         assistant_message = gr.Textbox(label="💫🌠Starling Assistant Message", lines=2)
         user_message = gr.Textbox(label="Your Message", lines=3)
     with gr.Row():
-        do_sample = gr.Checkbox(label="Advanced", value=True)
-    
+        mode = gr.Radio(choices=["Assistant", "Coder"], value="Assistant", label="Mode")
+        do_sample = gr.Checkbox(label="Advanced", value=True)    
     with gr.Accordion("Advanced Settings", open=lambda do_sample: do_sample):
         with gr.Row():
             temperature = gr.Slider(label="Temperature", value=0.4, minimum=0.05, maximum=1.0, step=0.05)
@@ -98,7 +101,7 @@ with gr.Blocks(theme="ParityError/Anime") as demo:
 
     submit_button.click(
         gradio_starling,
-        inputs=[user_message, assistant_message, system_prompt, do_sample, temperature, max_new_tokens, top_p, repetition_penalty],
+        inputs=[user_message, assistant_message, system_prompt, mode, do_sample, temperature, max_new_tokens, top_p, repetition_penalty],
         outputs=output_text
     )