Update app.py

app.py

@@ -1,31 +1,38 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from huggingface_hub import InferenceClient, TextGenerationStreamOutput
 import random
 from transformers import AutoTokenizer
 from mySystemPrompt import SYSTEM_PROMPT
+from datetime import datetime
 
 
 # Model which is used
 checkpoint = "CohereForAI/c4ai-command-r-plus"
+checkpoint = "mistral-community/Mixtral-8x22B-v0.1"
+path_to_log = "log.txt"
 # Inference client with the model (And HF-token if needed)
 client = InferenceClient(checkpoint)
 tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-# Tokenizer chat template correction(Only works for mistral models)
-#chat_template = open("mistral-instruct.jinja").read()
-#chat_template = chat_template.replace('    ', '').replace('\n', '')
-#tokenizer.chat_template = chat_template
-
-def format_prompt(message,chatbot,system_prompt):
-    messages = [{"role": "system","content": system_prompt}]
-    for user_message,bot_message in chatbot:
-        messages.append({"role": "user", "content":user_message})
-        messages.append({"role": "assistant", "content":bot_message})
-    messages.append({"role": "user", "content":message})
-    newPrompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, return_tensors="pt")
-    print(newPrompt)
-    #newPrompt = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
+
+if checkpoint == "mistral-community/Mixtral-8x22B-v0.1":
+    # Tokenizer chat template correction(Only works for mistral models)
+    chat_template = open("mistral-instruct.jinja").read()
+    chat_template = chat_template.replace('    ', '').replace('\n', '')
+    tokenizer.chat_template = chat_template
+
+
+def format_prompt(message, chatbot, system_prompt):
+    messages = [{"role": "system", "content": system_prompt}]
+    for user_message, bot_message in chatbot:
+        messages.append({"role": "user", "content": user_message})
+        messages.append({"role": "assistant", "content": bot_message})
+    messages.append({"role": "user", "content": message})
+    newPrompt = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True, return_tensors="pt")
+    # newPrompt = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
     return newPrompt
 
+
 def inference(message, history, systemPrompt=SYSTEM_PROMPT, temperature=0.9, maxTokens=512, topP=0.9, repPenalty=1.1):
     # Updating the settings for the generation
     client_settings = dict(
@@ -41,16 +48,22 @@ def inference(message, history, systemPrompt=SYSTEM_PROMPT, temperature=0.9, max
     )
     # Generating the response by passing the prompt in right format plus the client settings
     stream = client.text_generation(format_prompt(message, history, systemPrompt),
-                            **client_settings)
+                                    **client_settings)
     # Reading the stream
     partial_response = ""
     for stream_part in stream:
-        partial_response += stream_part.token.text
-        #yield partial_response
-    partial_response= partial_response.removesuffix("<|END_OF_TURN_TOKEN|>")
-    return partial_response
-    
-    
+        if not stream_part.token.special:
+            partial_response += stream_part.token.text
+            yield partial_response
+
+
+def event_voting(vote_data: gr.LikeData):
+    if vote_data.liked:
+        pass
+    else:
+        vote_message = f'{datetime.now().strftime(" %I: %M %p on %B %d, %Y")}:\n{vote_data.value}\n'
+        with open(path_to_log, "a") as f:
+            f.write(vote_message)
 
 
 myAdditionalInputs = [
@@ -105,6 +118,7 @@ myChatbot = gr.Chatbot(avatar_images=["./ava_m.png", "./avatar_franzi.jpg"],
                        show_copy_button=False,
                        likeable=True)
 
+
 myTextInput = gr.Textbox(lines=2,
                          max_lines=2,
                          placeholder="Send a message",
@@ -129,16 +143,19 @@ myClearButton = gr.Button(value="CLEAR",
                           size="sm")
 
 
-gr.ChatInterface(
-    inference,
-    chatbot=myChatbot,
-    textbox=myTextInput,
-    title="FRANZI-Bot 2.0",
-    theme=myTheme,
-    #additional_inputs=myAdditionalInputs,
-    submit_btn=mySubmitButton,
-    stop_btn="STOP",
-    retry_btn=myRetryButton,
-    undo_btn=myUndoButton,
-    clear_btn=myClearButton,
-).queue().launch(show_api=False)
+with gr.ChatInterface(inference,
+                      chatbot=myChatbot,
+                      textbox=myTextInput,
+                      title="FRANZI-Bot 2.0",
+                      theme=myTheme,
+                      # additional_inputs=myAdditionalInputs,
+                      submit_btn=mySubmitButton,
+                      stop_btn="STOP",
+                      retry_btn=myRetryButton,
+                      undo_btn=myUndoButton,
+                      clear_btn=myClearButton) as chatApp:
+    myChatbot.like(event_voting, None, None)
+
+
+if __name__ == "__main__":
+    chatApp.queue().launch(show_api=False)