Update app.py

app.py

@@ -87,98 +87,50 @@ model, tokenizer, description = setup(
     description=DESCRIPTION
 )
 
-
-def preprocess_conversation(
-    message: str, 
-    chat_history: list, 
-    system_prompt: str
-):
+def format_prompt(
+    message, 
+    history
+) -> str:
     """
-    Preprocess the conversation history by formatting it appropriately.
+    Format a prompt for dialogue generation using historical conversation data.
 
     Parameters
     ----------
     message : str
-        The user's message.
-    
-    chat_history : list
-        The conversation history, where each element is a tuple (user_message, assistant_response).
-    
-    system_prompt : str
-        The system prompt.
+        The user's current message or prompt.
+        
+    history : list of tuple
+        A list of tuples representing past interactions, where each tuple 
+        contains a user prompt and a corresponding bot response.
 
     Returns
     -------
-    list
-        The formatted conversation history.
-    """
-    conversation = []
-
-    if system_prompt:
-        conversation.append(
-            {
-                "role": "system", 
-                "content": system_prompt
-            }
-        )
-
-    for user, assistant in chat_history:
-        conversation.extend(
-            [
-                {
-                    "role": "user",
-                    "content": user
-                }, 
-                {
-                    "role": "assistant", 
-                    "content": assistant
-                }
-            ]
-        )
-
-    conversation.append(
-        {
-            "role": "user", 
-            "content": message
-        }
-    )
-
-    return conversation
-
-
-def trim_input_ids(
-    input_ids,
-    max_length
-):
+    str
+        Formatted prompt for dialogue generation, including the user's current 
+        message and historical conversation data.
+
+    Examples
+    --------
+    >>> message = "How are you?"
+    >>> history = [("Hi there!", "Hello!"), ("What's up?", "Not much.")]
+    >>> format_prompt(message, history)
+    '<s>[INST] Hi there! [/INST] Hello!</s> <s>[INST] What\'s up? [/INST] Not much.</s> <s>[INST] How are you? [/INST]'
     """
-    Trim the input token IDs if they exceed the maximum length.
-
-    Parameters
-    ----------
-    input_ids : torch.Tensor
-        The input token IDs.
+    prompt = "<s>"
     
-    max_length : int
-        The maximum length allowed.
+    for user_prompt, bot_response in history:
+        prompt += f"[INST] {user_prompt} [/INST]"
+        prompt += f" {bot_response}</s> "
+    prompt += f"[INST] {message} [/INST]"
 
-    Returns
-    -------
-    torch.Tensor
-        The trimmed input token IDs.
-    """
-    if input_ids.shape[1] > max_length:
-        input_ids = input_ids[:, -max_length:]
-        print(f"Trimmed input from conversation as it was longer than {max_length} tokens.")
-
-    return input_ids
+    return prompt
 
 
 @spaces.GPU
 def generate(
     message: str,
     chat_history: list,
-    system_prompt: str,
-    max_new_tokens: int = 1024,
+    max_new_tokens: int = 2048,
     temperature: float = 0.6,
     top_p: float = 0.9,
     top_k: int = 50,
@@ -196,9 +148,6 @@ def generate(
 
     chat_history : list
         A list containing tuples representing the conversation history. Each tuple should consist of two elements: the user's message and the assistant's response.
-   
-    system_prompt : str
-        The system prompt, if any, to be included in the conversation context.
     
     max_new_tokens : int, optional
         The maximum number of tokens to generate for the response (default is 1024).
@@ -228,10 +177,9 @@ def generate(
     global tokenizer
     global model
     
-    conversation = preprocess_conversation(
+    conversation = format_prompt(
         message=message, 
-        chat_history=chat_history, 
-        system_prompt=system_prompt
+        history=history
     )
 
     input_ids = tokenizer.apply_chat_template(
@@ -239,10 +187,6 @@ def generate(
         return_tensors="pt", 
         add_generation_prompt=True
     )
-    input_ids = trim_input_ids(
-        input_ids=input_ids, 
-        max_length=MAX_INPUT_TOKEN_LENGTH
-    )
 
     input_ids = input_ids.to(
         torch.device("cuda")
@@ -279,10 +223,16 @@ def generate(
 
     return "".join(outputs)
 
+
+chatbot = gr.Chatbot(
+    height=400,
+    show_copy_button=True
+)
+
 chat_interface = gr.ChatInterface(
     fn=generate,
+    chatbot=chatbot,
     additional_inputs=[
-        gr.Textbox(label="System prompt", lines=6),
         gr.Slider(
             label="Max new tokens",
             minimum=1,
@@ -314,9 +264,9 @@ chat_interface = gr.ChatInterface(
     ],
     fill_height=True,
     examples=[
-        ["implement snake game using pygame"],
+        ["Implement snake game using pygame"],
         ["Can you explain briefly to me what is the Python programming language?"],
-        ["write a program to find the factorial of a number"],
+        ["Write a program to find the factorial of a number"],
     ],
 )