songtao
/
entity_extraction
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

增加超时处理

main
million 2024-11-22 17:28:22 +08:00
parent 4c8228350b
commit 016e50e0d2
3 changed files with 8 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
config.py
View File

@ -2,3 +2,4 @@ API_KEY = "ollama"
MODEL = "qwen2.5:32b-instruct-q5_K_M" MODEL = "qwen2.5:32b-instruct-q5_K_M"
# MODEL = "qwen2.5:7b" # MODEL = "qwen2.5:7b"
LLM_BASE_URL = "http://localhost:11434/v1" LLM_BASE_URL = "http://localhost:11434/v1"
TIMEOUT_SECONDS = 300

6
entity_extraction_process.py
View File

@ -30,7 +30,11 @@ class entityExtractionProcess:
response = ( response = (
llm.chat.completions.create( llm.chat.completions.create(
model=config.MODEL, messages=messages, temperature=0, max_tokens=128_000 model=config.MODEL,
messages=messages,
temperature=0,
max_tokens=128_000,
timeout=config.TIMEOUT_SECONDS,
) )
.choices[0] .choices[0]
.message.content .message.content

4
main.py
View File

@ -20,12 +20,12 @@ DEBUG = False
class Main(QueueProcessor): class Main(QueueProcessor):
def processor_handle(self, input): def processor_handle(self, input):
print("input:", input) # print("input:", input)
currentEngineId = input["currentEngineId"] currentEngineId = input["currentEngineId"]
inputData = json.loads(input["inputData"]) inputData = json.loads(input["inputData"])
match currentEngineId: match currentEngineId:
case 3000: # InfoExtractor 实体信息提取 case 3000: # InfoExtractor 实体信息提取
logger.warning(f"get queueing inputData= {inputData}") # logger.warning(f"get queueing inputData= {inputData}")
input_text = inputData["text"] input_text = inputData["text"]
result = entityExtractionProcess.entity_extract(input_text) result = entityExtractionProcess.entity_extract(input_text)