70 lines
1.6 KiB
Python
70 lines
1.6 KiB
Python
from dotenv import load_dotenv
|
|
load_dotenv()
|
|
import time
|
|
from openai import OpenAI
|
|
from datetime import datetime
|
|
|
|
import config, debug, prompt
|
|
|
|
DEBUG = True
|
|
|
|
'''
|
|
备注：
|
|
目前使用ollama部署的qwen2.5:32b-instruct-q5_K_M,
|
|
如果输入字符串过长(大致超过1000个字符或1500token)会因上下文窗口过短导致提示词被遗忘，
|
|
进而输出不合规定的结果！
|
|
'''
|
|
|
|
llm = OpenAI(
|
|
base_url=config.LLM_BASE_URL,
|
|
api_key=config.API_KEY
|
|
)
|
|
|
|
def get_task() -> str:
|
|
# 从引擎中获取任务的接口
|
|
# 返回的类型应为字符串或None
|
|
raise NotImplementedError
|
|
|
|
def push_result(result):
|
|
# 将提取结果返回给引擎的接口
|
|
raise NotImplementedError
|
|
|
|
def entity_extract(input_text):
|
|
|
|
messages = [
|
|
{"role": "system", "content": prompt.ENTITY_EXTRACT},
|
|
{"role": "system", "content": f"今天的日期是：{str(datetime.today())}"},
|
|
{"role": "user", "content": f"{input_text}"}
|
|
]
|
|
|
|
timenow = time.time()
|
|
response = llm.chat.completions.create(
|
|
model=config.MODEL,
|
|
messages=messages,
|
|
temperature=0,
|
|
max_tokens=128_000
|
|
).choices[0].message.content
|
|
print(f"本次输出花费时间：{time.time() - timenow} 秒")
|
|
|
|
return response
|
|
|
|
def run():
|
|
if DEBUG:
|
|
input_text = debug.input_text
|
|
else:
|
|
input_text = get_task()
|
|
if not input_text:
|
|
return
|
|
response = entity_extract(input_text)
|
|
if DEBUG:
|
|
print(response)
|
|
else:
|
|
push_result(response)
|
|
|
|
if __name__ == "__main__":
|
|
if DEBUG:
|
|
run()
|
|
exit()
|
|
while True:
|
|
run()
|
|
time.sleep(5) |