Update app.py

app.py

@@ -9,6 +9,17 @@ model = AutoModelForCausalLM.from_pretrained(model_path)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
+def build_context(system_message, conversation_history, user_message):
+    context = f"<|start|>System:<|message|>{system_message}<|end|>\n"
+    if conversation_history:
+        for message in conversation_history:
+            if message["role"] == "user":
+                context += f"<|start|>User:<|message|>{message['content']}<|end|>\n"
+            elif message["role"] == "assistant":
+                context += f"<|start|>Assistant:<|message|>{message['content']}<|end|>\n"
+    context += f"<|start|>User:<|message|>{user_message}<|end|>\n<|start|>Assistant:<|message|>"
+    return context
+
 def generate_response(
     prompt,
     system_message,
@@ -19,25 +30,7 @@ def generate_response(
     repetition_penalty=1.031,
     top_k=55,
 ):
-    context = ""
-    if conversation_history:
-        recent = conversation_history[-30:] if len(conversation_history) > 30 else conversation_history
-        is_first_message = False
-        for i, message in enumerate(recent):
-            if i == 0:
-                is_first_message = True
-                context += (
-                    f"<|start|>User:<|message|>{system_message}<|end|>\n"
-                    f"<|start|>Assistant:<|message|>Hello, nice to meet you!<|end|>\n"
-                )
-            if message["role"] == "user":
-                context += f"<|start|>User:<|message|>{message['content']}<|end|>\n"
-            else:
-                context += f"<|start|>Assistant:<|message|>{message['content']}<|end|>\n"
-
-    formatted_input = (
-        f"{context}<|start|>User:<|message|>{prompt}<|end|>\n<|start|>Assistant:<|message|>"
-    )
+    formatted_input = build_context(system_message, conversation_history, prompt)
 
     inputs = tokenizer(
         formatted_input,
@@ -64,6 +57,9 @@ def generate_response(
     new_tokens = outputs[0][inputs.input_ids.shape[-1]:]
     response = tokenizer.decode(new_tokens, skip_special_tokens=False)
 
+    if "<|end|>" in response:
+        response = response.split("<|end|>")[0]
+
     return response.strip()
 
 def respond(
@@ -76,22 +72,17 @@ def respond(
     repetition_penalty,
     top_k,
 ):
-    conversation_history = history
     response = generate_response(
         message,
         system_message,
-        conversation_history,
+        history,
         max_tokens=max_tokens,
         temperature=temperature,
         top_p=top_p,
         repetition_penalty=repetition_penalty,
         top_k=top_k,
     )
-
-    if "<|end|>" in response:
-        response = response.split("<|end|>")[0]
-
-    return response.strip()
+    return response
 
 chatbot = gr.ChatInterface(
     respond,