human/test/test_nlp_only.py

#encoding = utf8
import json
import logging
import os
import time
import requests

from nlp import PunctuationSplit, DouBao, NLPCallback
from utils import config_logging

logger = logging.getLogger(__name__)
current_file_path = os.path.dirname(os.path.abspath(__file__))


# 接入点和apiKey生成、模型选择部分请访问豆包官方文档，本文仅给出请求体和访问路径
# 豆包大模型接口官方文档 https://console.volcengine.com/ark/region:ark+cn-beijing/endpoint/detail?Id=ep-20240826182225-kp7rp&Tab=api
# 部分内容参考通义千问文档实现 https://help.aliyun.com/zh/model-studio/developer-reference/use-qwen-by-calling-api?spm=a2c4g.11186623.0.0.28b919a1NbCP4i#e7932c7e33gvv


def __request(key, question):
    url = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"
    headers = {
        "Authorization": "Bearer " + key,
        "Content-Type": "application/json"
    }
    data = {
        "model": "ep-20241008152048-fsgzf",
        "messages": question,
        'stream': True
    }

    response = requests.post(url, headers=headers, json=data, stream=True)
    return response


class DisplayNLP(NLPCallback):
    def on_message(self, txt: str):
        print(txt)


def main():
    # 你的API_KEY
    # api_key = "c9635f9e-0f9e-4ca1-ac90-8af25a541b74"
    # __token = 'c9635f9e-0f9e-4ca1-ac90-8af25a541b74'
    # # 问题列表
    # msg_list = [
    #     {"role": "system", "content": "你是测试客服，是由字节跳动开发的 AI 人工智能助手"},
    #     {"role": "user", "content": "你好"}
    # ]
    # t = time.time()
    # response = __request(__token, msg_list)
    # if response.status_code != 200:
    #     print(f"请求失败，状态码：{response.status_code}")
    #     return
    #
    # for chunk in response.iter_lines():
    #     content = chunk.decode("utf-8").strip()
    #     if len(content) < 1:
    #         continue
    #
    #     content = content[5:]
    #     content = json.loads(content)
    #     print(f'-------dou_bao ask time:{time.time() - t:.4f}s, response:{content["choices"][0]["delta"]["content"]}')
    #
    # print("文件下载完成")

    split = PunctuationSplit()
    nlp = DouBao(None, split, DisplayNLP())
    nlp.ask('你好')
    nlp.ask('你是谁')
    nlp.ask('能做什么')
    time.sleep(5)
    nlp.stop()
    print("stop")


if __name__ == "__main__":
    config_logging('../logs/info.log', logging.INFO, logging.INFO)
    try:
        main()
    except KeyboardInterrupt:
        print("\nCaught Ctrl + C. Exiting")
    except Exception as e:
        print(e)