Update app.py

app.py

@@ -87,43 +87,77 @@ model, tokenizer, description = setup(
     description=DESCRIPTION
 )
 
-def format_prompt(
-    message, 
-    history
-) -> str:
+def preprocess_conversation(
+    message: str, 
+    history: list, 
+):
     """
-    Format a prompt for dialogue generation using historical conversation data.
+    Preprocess the conversation history by formatting it appropriately.
 
     Parameters
     ----------
     message : str
-        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.
+        The user's message.
+    
+    history : list
+        The conversation history, where each element is a tuple (user_message, assistant_response).
 
     Returns
     -------
-    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]'
+    list
+        The formatted conversation history.
     """
-    prompt = "<s>"
+    conversation = []
+
+    for user, assistant in 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
+):
+    """
+    Trim the input token IDs if they exceed the maximum length.
+
+    Parameters
+    ----------
+    input_ids : torch.Tensor
+        The input token IDs.
     
-    for user_prompt, bot_response in history:
-        prompt += f"[INST] {user_prompt} [/INST]"
-        prompt += f" {bot_response}</s> "
-    prompt += f"[INST] {message} [/INST]"
+    max_length : int
+        The maximum length allowed.
 
-    return prompt
+    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
 
 
 @spaces.GPU
@@ -177,9 +211,9 @@ def generate(
     global tokenizer
     global model
     
-    conversation = format_prompt(
+    conversation = preprocess_conversation(
         message=message, 
-        history=history
+        history=history, 
     )
 
     input_ids = tokenizer.apply_chat_template(
@@ -187,6 +221,11 @@ 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")
@@ -214,6 +253,7 @@ def generate(
         target=model.generate, 
         kwargs=generate_kwargs
     )
+    
     t.start()
 
     outputs = []