AI/dify
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git synced 2025-05-15 15:08:16 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: xinference last token being ignored (#1013)

Browse Source
This commit is contained in:
Uranus 2023-08-25 18:15:05 +08:00 committed by GitHub
parent 915e26527b
commit 2d9616c29c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 55 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

50
api/core/third_party/langchain/llms/xinference_llm.py vendored
View File

@ -3,8 +3,11 @@ from typing import Optional, List, Any, Union, Generator
from langchain.callbacks.manager import CallbackManagerForLLMRun from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.llms import Xinference from langchain.llms import Xinference
from langchain.llms.utils import enforce_stop_tokens from langchain.llms.utils import enforce_stop_tokens
from xinference.client import RESTfulChatglmCppChatModelHandle, \ from xinference.client import (
RESTfulChatModelHandle, RESTfulGenerateModelHandle RESTfulChatglmCppChatModelHandle,
RESTfulChatModelHandle,
RESTfulGenerateModelHandle,
)
class XinferenceLLM(Xinference): class XinferenceLLM(Xinference):
@ -29,7 +32,9 @@ class XinferenceLLM(Xinference):
model = self.client.get_model(self.model_uid) model = self.client.get_model(self.model_uid)
if isinstance(model, RESTfulChatModelHandle): if isinstance(model, RESTfulChatModelHandle):
generate_config: "LlamaCppGenerateConfig" = kwargs.get("generate_config", {}) generate_config: "LlamaCppGenerateConfig" = kwargs.get(
"generate_config", {}
)
if stop: if stop:
generate_config["stop"] = stop generate_config["stop"] = stop
@ -48,7 +53,9 @@ class XinferenceLLM(Xinference):
completion = model.chat(prompt=prompt, generate_config=generate_config) completion = model.chat(prompt=prompt, generate_config=generate_config)
return completion["choices"][0]["message"]["content"] return completion["choices"][0]["message"]["content"]
elif isinstance(model, RESTfulGenerateModelHandle): elif isinstance(model, RESTfulGenerateModelHandle):
generate_config: "LlamaCppGenerateConfig" = kwargs.get("generate_config", {}) generate_config: "LlamaCppGenerateConfig" = kwargs.get(
"generate_config", {}
)
if stop: if stop:
generate_config["stop"] = stop generate_config["stop"] = stop
@ -65,10 +72,14 @@ class XinferenceLLM(Xinference):
return combined_text_output return combined_text_output
else: else:
completion = model.generate(prompt=prompt, generate_config=generate_config) completion = model.generate(
prompt=prompt, generate_config=generate_config
)
return completion["choices"][0]["text"] return completion["choices"][0]["text"]
elif isinstance(model, RESTfulChatglmCppChatModelHandle): elif isinstance(model, RESTfulChatglmCppChatModelHandle):
generate_config: "ChatglmCppGenerateConfig" = kwargs.get("generate_config", {}) generate_config: "ChatglmCppGenerateConfig" = kwargs.get(
"generate_config", {}
)
if generate_config and generate_config.get("stream"): if generate_config and generate_config.get("stream"):
combined_text_output = "" combined_text_output = ""
@ -91,11 +102,20 @@ class XinferenceLLM(Xinference):
def _stream_generate( def _stream_generate(
self, self,
model: Union["RESTfulGenerateModelHandle", "RESTfulChatModelHandle", "RESTfulChatglmCppChatModelHandle"], model: Union[
"RESTfulGenerateModelHandle",
"RESTfulChatModelHandle",
"RESTfulChatglmCppChatModelHandle",
],
prompt: str, prompt: str,
run_manager: Optional[CallbackManagerForLLMRun] = None, run_manager: Optional[CallbackManagerForLLMRun] = None,
generate_config: Optional[ generate_config: Optional[
Union["LlamaCppGenerateConfig", "PytorchGenerateConfig", "ChatglmCppGenerateConfig"]] = None, Union[
"LlamaCppGenerateConfig",
"PytorchGenerateConfig",
"ChatglmCppGenerateConfig",
]
] = None,
) -> Generator[str, None, None]: ) -> Generator[str, None, None]:
""" """
Args: Args:
@ -108,7 +128,9 @@ class XinferenceLLM(Xinference):
Yields: Yields:
A string token. A string token.
""" """
if isinstance(model, (RESTfulChatModelHandle, RESTfulChatglmCppChatModelHandle)): if isinstance(
model, (RESTfulChatModelHandle, RESTfulChatglmCppChatModelHandle)
):
streaming_response = model.chat( streaming_response = model.chat(
prompt=prompt, generate_config=generate_config prompt=prompt, generate_config=generate_config
) )
@ -123,14 +145,10 @@ class XinferenceLLM(Xinference):
if choices: if choices:
choice = choices[0] choice = choices[0]
if isinstance(choice, dict): if isinstance(choice, dict):
if 'finish_reason' in choice and choice['finish_reason'] \ if "text" in choice:
and choice['finish_reason'] in ['stop', 'length']:
break
if 'text' in choice:
token = choice.get("text", "") token = choice.get("text", "")
elif 'delta' in choice and 'content' in choice['delta']: elif "delta" in choice and "content" in choice["delta"]:
token = choice.get('delta').get('content') token = choice.get("delta").get("content")
else: else:
continue continue
log_probs = choice.get("logprobs") log_probs = choice.get("logprobs")