fix: prompt for baichuan text generation models (#1299)

2023-10-10 13:01:18 +08:00 · 2023-10-10 13:01:18 +08:00 · 8480b0197b
commit 8480b0197b
parent df07fb5951
1 changed files with 6 additions and 0 deletions
--- a/api/core/model_providers/models/llm/baichuan_model.py
+++ b/api/core/model_providers/models/llm/baichuan_model.py
@ -37,6 +37,12 @@ class BaichuanModel(BaseLLM):
        prompts = self._get_prompt_from_messages(messages)
        return self._client.generate([prompts], stop, callbacks)

+    def prompt_file_name(self, mode: str) -> str:
+        if mode == 'completion':
+            return 'baichuan_completion'
+        else:
+            return 'baichuan_chat'
+
    def get_num_tokens(self, messages: List[PromptMessage]) -> int:
        """
        get num tokens of prompt messages.