Update app.py

app.py

@@ -10,23 +10,37 @@ def num_tokens_from_messages(messages, model="gpt-3.5-turbo-0613"):
     except KeyError:
         print("Warning: model not found. Using cl100k_base encoding.")
         encoding = tiktoken.get_encoding("cl100k_base")
-
-    # Define token counts based on model
-    tokens_per_message = 3
-    tokens_per_name = 1
-    if model == "gpt-3.5-turbo-0301":
-        tokens_per_message = 4  # {role/name}\n{content}\n
-        tokens_per_name = -1    # if there's a name, the role is omitted
+    if model in {
+        "gpt-3.5-turbo-0613",
+        "gpt-3.5-turbo-16k-0613",
+        "gpt-4-0314",
+        "gpt-4-32k-0314",
+        "gpt-4-0613",
+        "gpt-4-32k-0613",
+        }:
+        tokens_per_message = 3
+        tokens_per_name = 1
+    elif model == "gpt-3.5-turbo-0301":
+        tokens_per_message = 4  # every message follows <|start|>{role/name}\n{content}<|end|>\n
+        tokens_per_name = -1  # if there's a name, the role is omitted
     elif "gpt-3.5-turbo" in model:
-        print("Warning: gpt-3.5-turbo may update over time. Using default for gpt-3.5-turbo-0613.")
+        print("Warning: gpt-3.5-turbo may update over time. Returning num tokens assuming gpt-3.5-turbo-0613.")
+        return num_tokens_from_messages(messages, model="gpt-3.5-turbo-0613")
     elif "gpt-4" in model:
-        print("Warning: gpt-4 may update over time. Using default for gpt-4-0613.")
-
-    num_tokens = sum(tokens_per_message + sum(len(encoding.encode(value)) for key, value in message.items() if key != "name")
-                     + (len(encoding.encode(message.get("name", ""))) + tokens_per_name if "name" in message else 0)
-                     for message in messages)
-    num_tokens += 3  # every reply is primed with assistant
-
+        print("Warning: gpt-4 may update over time. Returning num tokens assuming gpt-4-0613.")
+        return num_tokens_from_messages(messages, model="gpt-4-0613")
+    else:
+        raise NotImplementedError(
+            f"""num_tokens_from_messages() is not implemented for model {model}."""
+        )
+    num_tokens = 0
+    for message in messages:
+        num_tokens += tokens_per_message
+        for key, value in message.items():
+            num_tokens += len(encoding.encode(value))
+            if key == "name":
+                num_tokens += tokens_per_name
+    num_tokens += 3  # every reply is primed with <|start|>assistant<|message|>
     return num_tokens
 
 
@@ -187,8 +201,8 @@ with app:
                 message_btn = gr.Button("发送", variant="primary", elem_classes="custom-button")
                 clear_btn = gr.Button("清除", variant="primary", elem_classes="custom-button")
         with gr.Column():
-            output_textbox = gr.Textbox(label="显示问题答案", lines=15, max_lines=15, elem_classes="custom-textbox",interactive=True,show_label=True,show_copy_button=True)
-            output_textbox2 = gr.Textbox(label="显示tokens数", lines=1, max_lines=2, elem_classes="custom-textbox",interactive=False,show_label=True)
+            output_textbox = gr.Textbox(label="显示问题答案", lines=14, max_lines=14, elem_classes="custom-textbox",interactive=True,show_label=True,show_copy_button=True)
+            output_textbox2 = gr.Textbox(label="显示tokens数", lines=1, max_lines=1, elem_classes="custom-textbox",interactive=False,show_label=True)
         message_btn.click(generate_response,inputs=[question,model_name,endpoint,transmit_api_key, official_api_key],outputs=output_textbox,show_progress=False) 
         clear_btn.click(clear_output,outputs=[question,output_textbox],show_progress=False)
         output_textbox.change(fn=num_tokens_from_messages,inputs=output_textbox,outputs=output_textbox2)