Update app.py

app.py

@@ -1,75 +1,73 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
-from transformers import AutoTokenizer  # Import the tokenizer
+from transformers import AutoTokenizer
+from langchain.memory import ConversationBufferWindowMemory
+from langchain.schema import HumanMessage, AIMessage, SystemMessage
 
-# Use the appropriate tokenizer for your model.
+# Initialize tokenizer and inference client
 tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta")
 client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
-# Define a maximum context length (tokens).  Check your model's documentation!
-MAX_CONTEXT_LENGTH = 4096  # Example: Adjust this based on your model!
+MAX_CONTEXT_LENGTH = 4096
 
-# Read the default prompt from a file
+# Load prompt from file
 with open("prompt.txt", "r") as file:
     nvc_prompt_template = file.read()
 
+# Initialize LangChain Memory (buffer window to keep recent conversation)
+memory = ConversationBufferWindowMemory(k=10, return_messages=True)
+
 def count_tokens(text: str) -> int:
-    """Counts the number of tokens in a given string."""
     return len(tokenizer.encode(text))
 
-def truncate_history(history: list[tuple[str, str]], system_message: str, max_length: int) -> list[tuple[str, str]]:
-    """Truncates the conversation history to fit within the maximum token limit.
-    Args:
-        history: The conversation history (list of user/assistant tuples).
-        system_message: The system message.
-        max_length: The maximum number of tokens allowed.
-    Returns:
-        The truncated history.
-    """
-    truncated_history = []
-    system_message_tokens = count_tokens(system_message)
-    current_length = system_message_tokens
-
-    # Iterate backwards through the history (newest to oldest)
-    for user_msg, assistant_msg in reversed(history):
-        user_tokens = count_tokens(user_msg) if user_msg else 0
-        assistant_tokens = count_tokens(assistant_msg) if assistant_msg else 0
-        turn_tokens = user_tokens + assistant_tokens
-
-        if current_length + turn_tokens <= max_length:
-            truncated_history.insert(0, (user_msg, assistant_msg))  # Add to the beginning
-            current_length += turn_tokens
+def truncate_history(messages, max_length):
+    truncated_messages = []
+    total_tokens = 0
+
+    for message in reversed(messages):
+        message_tokens = count_tokens(message.content)
+        if total_tokens + message_tokens <= max_length:
+            truncated_messages.insert(0, message)
+            total_tokens += message_tokens
         else:
-            break  # Stop adding turns if we exceed the limit
+            break
 
-    return truncated_history
+    return truncated_messages
 
 def respond(
     message,
-    history: list[tuple[str, str]],
+    history,
     system_message,
     max_tokens,
     temperature,
     top_p,
 ):
-    """Responds to a user message, maintaining conversation history, using special tokens and message list."""
     formatted_system_message = nvc_prompt_template
 
-    truncated_history = truncate_history(history, formatted_system_message, MAX_CONTEXT_LENGTH - max_tokens - 100)  # Reserve space for the new message and some generation
+    # Retrieve conversation history from LangChain memory
+    memory.save_context({"input": message}, {"output": ""})
+    chat_history = memory.load_memory_variables({})["history"]
+
+    # Truncate history to ensure it fits within context window
+    max_history_tokens = MAX_CONTEXT_LENGTH - max_tokens - count_tokens(formatted_system_message) - 100
+    truncated_chat_history = truncate_history(chat_history, max_history_tokens)
 
-    messages = [{"role": "system", "content": formatted_system_message}]  # Start with system message
-    for user_msg, assistant_msg in truncated_history:
-        if user_msg:
-            messages.append({"role": "user", "content": f"<|user|>\n{user_msg}</s>"})
-        if assistant_msg:
-            messages.append({"role": "assistant", "content": f"<|assistant|>\n{assistant_msg}</s>"})
+    # Construct the messages for inference
+    messages = [SystemMessage(content=formatted_system_message)]
+    messages.extend(truncated_chat_history)
+    messages.append(HumanMessage(content=message))
 
-    messages.append({"role": "user", "content": f"<|user|>\n{message}</s>"})
+    # Convert LangChain messages to the format required by HuggingFace client
+    formatted_messages = []
+    for msg in messages:
+        role = "system" if isinstance(msg, SystemMessage) else "user" if isinstance(msg, HumanMessage) else "assistant"
+        content = f"<|{role}|>\n{msg.content}</s>"
+        formatted_messages.append({"role": role, "content": content})
 
     response = ""
     try:
         for chunk in client.chat_completion(
-            messages,
+            formatted_messages,
             max_tokens=max_tokens,
             stream=True,
             temperature=temperature,
@@ -78,6 +76,10 @@ def respond(
             token = chunk.choices[0].delta.content
             response += token
             yield response
+
+        # Save AI's response in LangChain memory
+        memory.chat_memory.add_ai_message(response)
+
     except Exception as e:
         print(f"An error occurred: {e}")
         yield "I'm sorry, I encountered an error. Please try again."