Skip to content

Commit

Permalink
chore: Fix baichuan API pylint error (#1010)
Browse files Browse the repository at this point in the history
Co-authored-by: csunny <[email protected]>
  • Loading branch information
fangyinc and csunny authored Jan 3, 2024
1 parent 9ad70a2 commit c871d65
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 6 deletions.
2 changes: 1 addition & 1 deletion dbgpt/_private/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,7 @@ def __init__(self) -> None:
# baichuan proxy
self.bc_proxy_api_key = os.getenv("BAICHUAN_PROXY_API_KEY")
self.bc_model_name = os.getenv("BAICHUN_MODEL_NAME", "Baichuan2-Turbo-192k")
if self.bc_proxy_api_key and self.bc_proxy_api_secret:
if self.bc_proxy_api_key and self.bc_model_name:
os.environ["bc_proxyllm_proxy_api_key"] = self.bc_proxy_api_key
os.environ["bc_proxyllm_proxyllm_backend"] = self.bc_model_name

Expand Down
17 changes: 12 additions & 5 deletions dbgpt/model/proxy/llms/baichuan.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,9 @@
BAICHUAN_DEFAULT_MODEL = "Baichuan2-Turbo-192k"


def baichuan_generate_stream(model: ProxyModel, tokenizer=None, params=None, device=None, context_len=4096):
def baichuan_generate_stream(
model: ProxyModel, tokenizer=None, params=None, device=None, context_len=4096
):
url = "https://api.baichuan-ai.com/v1/chat/completions"

model_params = model.get_params()
Expand Down Expand Up @@ -63,22 +65,27 @@ def baichuan_generate_stream(model: ProxyModel, tokenizer=None, params=None, dev
text += content
yield text


def main():
model_params = ProxyModelParameters(
model_name="not-used",
model_path="not-used",
proxy_server_url="not-used",
proxy_api_key="YOUR_BAICHUAN_API_KEY",
proxyllm_backend="Baichuan2-Turbo-192k"
proxyllm_backend="Baichuan2-Turbo-192k",
)
final_text = ""
for part in baichuan_generate_stream(
model=ProxyModel(model_params=model_params),
params={"messages": [ModelMessage(
role=ModelMessageRoleType.HUMAN,
content="背诵《论语》第一章")]}):
params={
"messages": [
ModelMessage(role=ModelMessageRoleType.HUMAN, content="背诵《论语》第一章")
]
},
):
final_text = part
print(final_text)


if __name__ == "__main__":
main()

0 comments on commit c871d65

Please sign in to comment.