Update app.py

app.py

@@ -10,14 +10,15 @@ openai.api_key = os.environ.get("OPENAI_API_KEY")
 
 BASE_SYSTEM_MESSAGE = """"""
 
-def make_prediction(prompt, max_tokens=None, temperature=None, top_p=None, top_k=None, repetition_penalty=None):
-    openai.api_base = os.environ.get("OPENAI_API_BASE")
-    completion1 = openai.Completion.create(model="wizardcoder-python-34b-v1.0.Q5_K_M.gguf", prompt=prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty, stream=True, stop=["</s>", "<|im_end|>"])
-    openai.api_base = os.environ.get("OPENAI_API_BASE2")
-    completion2 = openai.Completion.create(model="wizardcoder-python-34b-v1.0.Q5_K_M.gguf", prompt=prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty, stream=True, stop=["</s>", "<|im_end|>"])
-    for chunk in completion1:
-        yield chunk["choices"][0]["text"]
-    for chunk in completion2:
+def make_prediction(chat, prompt, max_tokens=None, temperature=None, top_p=None, top_k=None, repetition_penalty=None):
+    if chat="Chatbot1":
+        openai.api_base = os.environ.get("OPENAI_API_BASE")
+        completion = openai.Completion.create(model="wizardcoder-python-34b-v1.0.Q5_K_M.gguf", prompt=prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty, stream=True, stop=["</s>", "<|im_end|>"])
+    elif chat="Chatbot2":
+        openai.api_base = os.environ.get("OPENAI_API_BASE2")
+        completion = openai.Completion.create(model="wizardcoder-python-34b-v1.0.Q5_K_M.gguf", prompt=prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty, stream=True, stop=["</s>", "<|im_end|>"])
+
+    for chunk in completion:
         yield chunk["choices"][0]["text"]
 
 
@@ -45,7 +46,8 @@ def chat(history, system_message, max_tokens, temperature, top_p, top_k, repetit
     # remove last space from assistant, some models output a ZWSP if you leave a space
     messages = messages.rstrip()
 
-    prediction = make_prediction(
+    prediction1 = make_prediction(
+        chat="Chatbot1",
         messages,
         max_tokens=max_tokens,
         temperature=temperature,
@@ -53,7 +55,27 @@ def chat(history, system_message, max_tokens, temperature, top_p, top_k, repetit
         top_k=top_k,
         repetition_penalty=repetition_penalty,
     )
-    for tokens in prediction:
+    prediction2 = = make_prediction(
+        chat="Chatbot1",
+        messages,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        repetition_penalty=repetition_penalty,
+    )
+    for tokens in prediction1:
+        tokens = re.findall(r'(.*?)(\s|$)', tokens)
+        for subtoken in tokens:
+            subtoken = "".join(subtoken)
+            # Remove "Response\n" if it's at the beginning of the assistant's output
+            if subtoken.startswith("Response"):
+                subtoken = subtoken[len("Response"):]
+            answer = subtoken
+            history[-1][1] += answer
+            # stream the response
+            yield history, history, ""
+    for tokens in prediction2:
         tokens = re.findall(r'(.*?)(\s|$)', tokens)
         for subtoken in tokens:
             subtoken = "".join(subtoken)
@@ -82,7 +104,6 @@ with gr.Blocks(css=CSS) as demo:
         with gr.Column():
             gr.Markdown(f"""
                     ## This demo is an unquantized GPU chatbot of [WizardCoder-Python-34B-V1.0-GGUF](https://huggingface.co/TheBloke/WizardCoder-Python-34B-V1.0-GGUF)
-                    Brought to you by your friends at Alignment Lab AI, garage-bAInd, Open Access AI Collective, and OpenChat!
                     """)
     with gr.Row():
         gr.Markdown("# 🔍 WizardCoder-Python-34B-V1.0-GGUF Playground Space! 🔎")