Update app.py

app.py

@@ -24,70 +24,68 @@ def create_formatted_history(history_messages: List[dict]) -> List[Tuple[str, st
     formatted_history = []
     user_messages = []
     assistant_messages = []
-
     for message in history_messages:
         if message["role"] == "user":
             user_messages.append(message["content"])
         elif message["role"] == "assistant":
             assistant_messages.append(message["content"])
-
         if user_messages and assistant_messages:
             formatted_history.append(
                 ("".join(user_messages), "".join(assistant_messages))
             )
             user_messages = []
             assistant_messages = []
-
     # Append any remaining messages
     if user_messages:
         formatted_history.append(("".join(user_messages), None))
     elif assistant_messages:
         formatted_history.append((None, "".join(assistant_messages)))
-
     return formatted_history
 
-def chat(message: str, state: List[Dict[str, str]]) -> Generator[Tuple[List[Tuple[str, str]], List[Dict[str, str]]], None, None]:
-    history_messages = state
-    if history_messages == None:
-        history_messages = []
-        history_messages.append({"role": "system", "content": "A helpful assistant."})
-
-    history_messages.append({"role": "user", "content": message})
-    history_messages.append({"role": "assistant", "content": ""})
-
-    # Tokenize user input and prepare input tensor
-    input_ids = tokenizer.encode(message, return_tensors='pt').to(device)
-
-    if input_ids.size(-1) == 0:
-        response_message = "Input was empty after tokenization. Please try again."
-    else:
-        # Generate tokens one by one
-        with torch.no_grad():
-            for _ in range(100):  # Limit generation to 50 tokens
-                outputs = model(input_ids)
-                next_token_logits = outputs.logits[:, -1, :]
-                next_token_id = torch.argmax(next_token_logits, dim=-1)
-                input_ids = torch.cat([input_ids, next_token_id.unsqueeze(-1)], dim=-1)
-
-                # Decode and append the latest token
-                decoded_token = tokenizer.decode(next_token_id)
-                history_messages[-1]["content"] += decoded_token
-
-                # Stop if the model generates the end-of-sequence token
-                if next_token_id.item() == tokenizer.eos_token_id:
-                    break
-
-        response_message = history_messages[-1]["content"]
-
-    formatted_history = create_formatted_history(history_messages)
-    yield formatted_history, history_messages
+class ConversationHistory:
+  def __init__(self):
+    self.messages: List[Tuple[str, str]] = []  # Stores conversation history
+
+  def append(self, user_message: str, assistant_message: str):
+    self.messages.append((user_message, assistant_message))
+
+  def get_formatted_history(self):
+    return create_formatted_history(create_history_messages(self.messages))
+
+def chat(message: str, conversation_history: ConversationHistory) -> Generator[Tuple[List[Tuple[str, str]], ConversationHistory], None, None]:
+  # Update history
+  conversation_history.append(message, "")
+
+  # Tokenize user input and prepare input tensor
+  input_ids = tokenizer.encode(message, return_tensors='pt').to(device)
+  if input_ids.size(-1) == 0:
+      response_message = "Input was empty after tokenization. Please try again."
+  else:
+      # Generate tokens one by one
+      with torch.no_grad():
+          for _ in range(100):  # Limit generation to 50 tokens
+              outputs = model(input_ids)
+              next_token_logits = outputs.logits[:, -1, :]
+              next_token_id = torch.argmax(next_token_logits, dim=-1)
+              input_ids = torch.cat([input_ids, next_token_id.unsqueeze(-1)], dim=-1)
+              # Decode and append the latest token
+              decoded_token = tokenizer.decode(next_token_id)
+              conversation_history.messages[-1] = (conversation_history.messages[-1][0], decoded_token)
+              # Stop if the model generates the end-of-sequence token
+              if next_token_id.item() == tokenizer.eos_token_id:
+                  break
+          response_message = conversation_history.messages[-1][1]
+
+  # Yield formatted history and updated conversation history
+  yield conversation_history.get_formatted_history(), conversation_history
 
 chatbot = gr.Chatbot(label="Chat")
+conversation_history = ConversationHistory()  # Initialize conversation history
+
 iface = gr.Interface(
     fn=chat,
-    inputs=[gr.Textbox(placeholder="Hello! How are you?", label="Message"), "state"],
-    outputs=[chatbot, "state"],
+    inputs=[gr.Textbox(placeholder="Hello! How are you?", label="Message")],
+    outputs=[chatbot, conversation_history],
     allow_flagging="never",
 )
-
-iface.queue().launch()
+iface.queue().launch()