mirror of
https://github.com/JasonYANG170/CodeGeeX4.git
synced 2024-10-19 23:41:53 +00:00
36 lines
968 B
Python
36 lines
968 B
Python
|
import json
|
||
|
|
|
||
|
|
from backend.apis.api import API
|
||
|
|
from backend.apis.bing import BingSearchAPI
|
||
|
|
from backend.models.codegeex import model, tokenizer
|
||
|
|
from backend.utils.chat import build_model_input, SYS_PROMPT
|
||
|
|
|
||
|
|
|
||
|
|
def chat(query: str, history: list[list[str]] = None):
|
||
|
|
if not history:
|
||
|
|
history = []
|
||
|
|
|
||
|
|
ans = ""
|
||
|
|
|
||
|
|
# Search with bing
|
||
|
|
api: API = BingSearchAPI()
|
||
|
|
search_res = api.call(query=query, history=history)
|
||
|
|
ans += "搜索结果".center(100, "-") + '\n'
|
||
|
|
ans += "```json\n" + json.dumps(search_res, indent=4, ensure_ascii=False) + "\n```\n"
|
||
|
|
yield ans
|
||
|
|
|
||
|
|
# Build model's input
|
||
|
|
inputs: str = build_model_input(query, search_res)
|
||
|
|
|
||
|
|
# Generate response
|
||
|
|
ans += "模型回复".center(100, "-") + '\n'
|
||
|
|
yield ans
|
||
|
|
response, _ = model.chat(
|
||
|
|
tokenizer,
|
||
|
|
query=inputs,
|
||
|
|
history=[{"role": "system", "content": SYS_PROMPT}],
|
||
|
|
max_new_tokens=1024,
|
||
|
|
temperature=0.2
|
||
|
|
)
|
||
|
|
yield ans + response
|