vrr/kvcache-ai-ktransformers
2
0
Fork 0
mirror of https://github.com/kvcache-ai/ktransformers.git synced 2025-09-06 20:49:55 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Update completions.py

Browse source
This commit is contained in:
Yuhao Tsui 2025-03-26 15:39:46 +08:00 • committed by GitHub
parent 52fa671c10
commit 84164f584c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 12 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
ktransformers/server/api/ollama/completions.py
View file

@ -85,8 +85,12 @@ async def generate(request: Request, input: OllamaGenerateCompletionRequest):
return check_link_response(request, inner())
else:
complete_response = ""
async for token in interface.inference(input.prompt, id):
complete_response += token
async for res in interface.inference(input.prompt, id):
if isinstance(res, RawUsage):
raw_usage = res
else:
token, finish_reason = res
complete_response += token
response = OllamaGenerationResponse(
model=config.model_name,
created_at=str(datetime.now()),
@ -187,8 +191,12 @@ async def chat(request: Request, input: OllamaChatCompletionRequest):
complete_response = ""
eval_count = 0
async for token in interface.inference(prompt, id):
complete_response += token
async for res in interface.inference(prompt, id):
if isinstance(res, RawUsage):
raw_usage = res
else:
token, finish_reason = res
complete_response += token
eval_count += 1
end_time = time()