P2L Extra body for cost param (#3666)

lm-sys · Jan 23, 2025 · d161b64 · d161b64
1 parent 1ffd4a6
commit d161b64
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/fastchat/serve/api_provider.py b/fastchat/serve/api_provider.py
@@ -23,6 +23,7 @@ def get_api_provider_stream_iter(
     top_p,
     max_new_tokens,
     state,
+    extra_body=None,
 ):
     if model_api_dict["api_type"] == "openai":
         if model_api_dict.get("vision-arena", False):
@@ -256,6 +257,7 @@ def get_api_provider_stream_iter(
             max_new_tokens,
             api_base=model_api_dict["api_base"],
             api_key=model_api_dict["api_key"],
+            extra_body=extra_body,
         )
     else:
         raise NotImplementedError()
@@ -431,6 +433,7 @@ def p2l_api_stream_iter(
     max_new_tokens,
     api_base=None,
     api_key=None,
+    extra_body=None,
 ):
     import openai
 
@@ -459,6 +462,7 @@ def p2l_api_stream_iter(
         "temperature": None,
         "top_p": None,
         "max_new_tokens": max_new_tokens,
+        "extra_body": extra_body,
     }
     logger.info(f"==== request ====\n{gen_params}")
 
@@ -467,6 +471,7 @@ def p2l_api_stream_iter(
         messages=messages,
         max_tokens=max_new_tokens,
         stream=True,
+        extra_body=extra_body,
     )
     text = ""
     for chunk_idx, chunk in enumerate(res):

diff --git a/fastchat/serve/gradio_web_server.py b/fastchat/serve/gradio_web_server.py
@@ -526,6 +526,8 @@ def bot_response(
         if not custom_system_prompt:
             conv.set_system_message("")
 
+        extra_body = None
+
         if use_recommended_config:
             recommended_config = model_api_dict.get("recommended_config", None)
             if recommended_config is not None:
@@ -534,6 +536,7 @@ def bot_response(
                 max_new_tokens = recommended_config.get(
                     "max_new_tokens", max_new_tokens
                 )
+                extra_body = recommended_config.get("extra_body", None)
 
         stream_iter = get_api_provider_stream_iter(
             conv,
@@ -543,6 +546,7 @@ def bot_response(
             top_p,
             max_new_tokens,
             state,
+            extra_body=extra_body,
         )
 
     html_code = ' <span class="cursor"></span> '