Update app.py

app.py

@@ -12,12 +12,6 @@ tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)
 
 BASE_SYSTEM_MESSAGE = "I carefully provide accurate, factual, thoughtful, nuanced answers and am brilliant at reasoning."
 
-def make_prediction(prompt, max_tokens=None, temperature=None, top_p=None, top_k=None, repetition_penalty=None):
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
-    out = model.generate(input_ids, max_length=max_tokens, temperature=temperature, top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty)
-    text = tokenizer.decode(out[0], skip_special_tokens=True)
-    yield text
-
 def clear_chat(chat_history_state, chat_message):
     chat_history_state = []
     chat_message = ''
@@ -28,13 +22,21 @@ def user(message, history):
     history.append([message, ""])
     return "", history
 
-def regenerate(_chatbot, _task_history):
+def regenerate(_chatbot, _task_history, system_msg, max_tokens, temperature, top_p, top_k, repetition_penalty):
     if not _task_history:
         yield _chatbot
         return
-    item = _task_history.pop(-1)
-    _chatbot.pop(-1)
-    yield from make_prediction(item[0])  # Assuming make_prediction is the function you want to use for regeneration
+    # Remove the last bot message
+    last_message = _task_history[-1]
+    if last_message[1]:  # Check if the bot actually sent a message last
+        last_message[1] = ""
+    
+    # Regenerate the message using the chat function
+    new_history, _, _ = chat(_task_history, system_msg, max_tokens, temperature, top_p, top_k, repetition_penalty)
+    
+    # Update the chatbot state
+    yield new_history
+
 
 
 def chat(history, system_message, max_tokens, temperature, top_p, top_k, repetition_penalty):
@@ -66,7 +68,7 @@ assistant
     )
     
     # Decodificar a saída
-    decoded_output = tokenizer.decode(output[0])
+    decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)
     assistant_response = decoded_output.split('assistant')[-1].strip()  # Pegar apenas a última resposta do assistente
     
     # Atualizar o histórico
@@ -134,7 +136,7 @@ with gr.Blocks() as demo:
     stop.click(fn=None, inputs=None, outputs=None, cancels=[submit_click_event], queue=False)
 
     regen_btn.click(
-        fn=regenerate, inputs=[chatbot, chat_history_state], outputs=[chatbot], queue=True
+        fn=regenerate, inputs=[chatbot, chat_history_state, system_msg, max_tokens, temperature, top_p, top_k, repetition_penalty], outputs=[chatbot], queue=True
     )
 
 demo.queue(max_size=128, concurrency_count=2)