add oss madule
This commit is contained in:
parent
61df642a34
commit
447d37a274
89
oss.py
Normal file
89
oss.py
Normal file
|
|
@ -0,0 +1,89 @@
|
|||
async def single_async_item(
|
||||
api_url,
|
||||
api_key,
|
||||
item,
|
||||
reasoning_effort,
|
||||
temperature,
|
||||
top_p,
|
||||
semaphore_number,
|
||||
model_name,
|
||||
priority,
|
||||
output_schema=None,
|
||||
max_token=4096,
|
||||
print_logs=False,
|
||||
return_reason=False,
|
||||
stop=None,
|
||||
return_used_token=False,
|
||||
timeout=300,
|
||||
):
|
||||
try:
|
||||
async with AsyncOpenAI(
|
||||
base_url=api_url, api_key=api_key
|
||||
) as client:
|
||||
semaphore = asyncio.Semaphore(semaphore_number)
|
||||
async with semaphore:
|
||||
messages = [{"role": "user", "content": item["user_prompt"]}]
|
||||
if item.get("system_prompt"):
|
||||
messages.insert(
|
||||
0, {"role": "system", "content": item["system_prompt"]}
|
||||
)
|
||||
if item.get("assistant_prompt"):
|
||||
messages.append(
|
||||
{"role": "assistant", "content": item["assistant_prompt"]}
|
||||
)
|
||||
|
||||
coro = client.chat.completions.parse(
|
||||
model=model_name,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
max_tokens=max_token,
|
||||
stop=stop,
|
||||
response_format=output_schema,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body={"priority": priority},
|
||||
)
|
||||
response = await asyncio.wait_for(coro, timeout=timeout)
|
||||
|
||||
if print_logs:
|
||||
print(f"parse response ---- {response}")
|
||||
|
||||
parsed_obj = response.choices[0].message.parsed
|
||||
# print(f'parsed_obj {parsed_obj}')
|
||||
if parsed_obj is None:
|
||||
return {
|
||||
"error": "Failed to parse response",
|
||||
"raw": str(response),
|
||||
}
|
||||
|
||||
parsed_obj = output_schema.model_validate(parsed_obj)
|
||||
# Validate just in case (optional, چون .parse already does it)
|
||||
if return_reason:
|
||||
reasoning_content = response.choices[
|
||||
0
|
||||
].message.reasoning_content
|
||||
if return_used_token:
|
||||
_total_token = response.usage.total_tokens
|
||||
item["llm_output"] = (
|
||||
parsed_obj.model_dump(),
|
||||
str(reasoning_content),
|
||||
int(_total_token),
|
||||
)
|
||||
return item
|
||||
|
||||
item["llm_output"] = (
|
||||
parsed_obj.model_dump(),
|
||||
str(reasoning_content)
|
||||
)
|
||||
return item
|
||||
|
||||
item["llm_output"] = parsed_obj.model_dump()
|
||||
return item
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
print(f"⏳ Timeout on item {item}")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
print(f"⚠️ Error __process_item {item}: {traceback.print_exc()}")
|
||||
return None
|
||||
Loading…
Reference in New Issue
Block a user