Update app.py

app.py

@@ -30,15 +30,15 @@ model.eval()
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:50'
 
 class StarlingBot:
-    def __init__(self, system_prompt="I am Starling-7B by Tonic-AI, I ready to do anything to help my user."):
+    def __init__(self, assistant_message="I am Starling-7B by Tonic-AI, I am ready to do anything to help my user."):
         self.system_prompt = system_prompt
 
-    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):
+    def predict(self, user_message, assistant_message, mode, do_sample, temperature=0.4, max_new_tokens=700, top_p=0.99, repetition_penalty=1.9):
         try:
             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:"
+                conversation = f"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 ''}  Code User:: {user_message}  Code Assistant:"
+                conversation = f"Code Assistant: {assistant_message if assistant_message else ''}  Code 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(
@@ -72,10 +72,10 @@ examples = [
         1.9,  # repetition_penalty
     ]
 ]
-# Initialize StarlingBot
+
 starling_bot = StarlingBot()
 
-def gradio_starling(user_message, assistant_message, system_prompt, mode, do_sample, temperature, max_new_tokens, top_p, repetition_penalty):
+def gradio_starling(user_message, assistant_message, mode, do_sample, temperature, max_new_tokens, top_p, repetition_penalty):
     response = starling_bot.predict(user_message, assistant_message, system_prompt, mode, do_sample, temperature, max_new_tokens, top_p, repetition_penalty)
     return response
 
@@ -83,8 +83,7 @@ with gr.Blocks(theme="ParityError/Anime") as demo:
     gr.Markdown(title)  
     gr.Markdown(description)  
     with gr.Row():
-        system_prompt = gr.Textbox(label="Optional💫🌠Starling System Prompt", lines=2)
-        assistant_message = gr.Textbox(label="💫🌠Starling Assistant Message", lines=2)
+        assistant_message = gr.Textbox(label="Optional💫🌠Starling Assistant Message", lines=2)
         user_message = gr.Textbox(label="Your Message", lines=3)
     with gr.Row():
         mode = gr.Radio(choices=["Assistant", "Coder"], value="Assistant", label="Mode")
@@ -101,7 +100,7 @@ with gr.Blocks(theme="ParityError/Anime") as demo:
 
     submit_button.click(
         gradio_starling,
-        inputs=[user_message, assistant_message, system_prompt, mode, do_sample, temperature, max_new_tokens, top_p, repetition_penalty],
+        inputs=[user_message, assistant_message, mode, do_sample, temperature, max_new_tokens, top_p, repetition_penalty],
         outputs=output_text
     )