40 lines
1.1 KiB
Python
40 lines
1.1 KiB
Python
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
import time
|
|
from openai import OpenAI
|
|
from datetime import datetime
|
|
|
|
import config, debug, prompt, json
|
|
|
|
"""
|
|
备注：
|
|
目前使用ollama部署的qwen2.5:32b-instruct-q5_K_M,
|
|
如果输入字符串过长(大致超过1000个字符或1500token)会因上下文窗口过短导致提示词被遗忘，
|
|
进而输出不合规定的结果！
|
|
"""
|
|
|
|
llm = OpenAI(base_url=config.LLM_BASE_URL, api_key=config.API_KEY)
|
|
|
|
|
|
class entityExtractionProcess:
|
|
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 json.loads(response)
|