Repository: xorbitsai/inference

Issue:
BUG: PyTorch generate config `max_new_tokens` not compatible with RESTful API
Should be `max_tokens`.

Reference: https://platform.openai.com/docs/api-reference/completions


Discussion:

Resolving patch:
diff --git a/examples/AI_podcast_ZH.py b/examples/AI_podcast_ZH.py
--- a/examples/AI_podcast_ZH.py
+++ b/examples/AI_podcast_ZH.py
@@ -263,7 +263,7 @@ def _base_sanitize_generate_config() -> PytorchGenerateConfig:
     pytorch_generate_config = PytorchGenerateConfig()
     pytorch_generate_config.setdefault("temperature", 0.7)
     pytorch_generate_config.setdefault("repetition_penalty", 1.0)
-    pytorch_generate_config.setdefault("max_new_tokens", 512)
+    pytorch_generate_config.setdefault("max_tokens", 512)
     pytorch_generate_config.setdefault("stream_interval", 2)
     pytorch_generate_config["stop"] = [f" {username} "]
     return pytorch_generate_config
diff --git a/xinference/model/llm/pytorch/core.py b/xinference/model/llm/pytorch/core.py
--- a/xinference/model/llm/pytorch/core.py
+++ b/xinference/model/llm/pytorch/core.py
@@ -38,7 +38,7 @@ class PytorchGenerateConfig(TypedDict, total=False):
     top_p: float
     top_k: int
     stream: bool
-    max_new_tokens: int
+    max_tokens: int
     echo: bool
     stop: Optional[Union[str, List[str]]]
     stop_token_ids: Optional[Union[int, List[int]]]
@@ -97,7 +97,7 @@ def _sanitize_generate_config(
             pytorch_generate_config = PytorchGenerateConfig()
         pytorch_generate_config.setdefault("temperature", 0.7)
         pytorch_generate_config.setdefault("repetition_penalty", 1.0)
-        pytorch_generate_config.setdefault("max_new_tokens", 512)
+        pytorch_generate_config.setdefault("max_tokens", 512)
         pytorch_generate_config.setdefault("stream_interval", 2)
         pytorch_generate_config["model"] = self.model_uid
         return pytorch_generate_config
diff --git a/xinference/model/llm/pytorch/utils.py b/xinference/model/llm/pytorch/utils.py
--- a/xinference/model/llm/pytorch/utils.py
+++ b/xinference/model/llm/pytorch/utils.py
@@ -100,7 +100,7 @@ def generate_stream(
     repetition_penalty = float(generate_config.get("repetition_penalty", 1.0))
     top_p = float(generate_config.get("top_p", 1.0))
     top_k = int(generate_config.get("top_k", -1))  # -1 means disable
-    max_new_tokens = int(generate_config.get("max_new_tokens", 256))
+    max_new_tokens = int(generate_config.get("max_tokens", 256))
     echo = bool(generate_config.get("echo", False))
     stop_str = generate_config.get("stop", None)
     stop_token_ids = generate_config.get("stop_token_ids", None) or []
@@ -332,7 +332,7 @@ def generate_stream_falcon(
     repetition_penalty = float(generate_config.get("repetition_penalty", 1.0))
     top_p = float(generate_config.get("top_p", 1.0))
     top_k = int(generate_config.get("top_k", 50))  # -1 means disable
-    max_new_tokens = int(generate_config.get("max_new_tokens", 256))
+    max_new_tokens = int(generate_config.get("max_tokens", 256))
     echo = bool(generate_config.get("echo", False))
     stop_str = generate_config.get("stop", None)
     stop_token_ids = generate_config.get("stop_token_ids", None) or []
@@ -507,7 +507,7 @@ def generate_stream_chatglm(
     temperature = float(generate_config.get("temperature", 1.0))
     repetition_penalty = float(generate_config.get("repetition_penalty", 1.0))
     top_p = float(generate_config.get("top_p", 1.0))
-    max_new_tokens = int(generate_config.get("max_new_tokens", 256))
+    max_new_tokens = int(generate_config.get("max_tokens", 256))
     echo = generate_config.get("echo", False)
 
     inputs = tokenizer([prompt], return_tensors="pt").to(model.device)


Test patch: