adapter.py 50.3 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
"""
Copyright 2023-2024 SGLang Team
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

16
"""Conversion between OpenAI APIs and native SRT APIs"""
Liangsheng Yin's avatar
Liangsheng Yin committed
17

18
import asyncio
19
import json
20
import logging
21
import os
22
23
import time
import uuid
24
from http import HTTPStatus
25
from typing import Dict, List
26

27
from fastapi import HTTPException, Request, UploadFile
28
from fastapi.responses import JSONResponse, StreamingResponse
29
from pydantic import ValidationError
30

31
32
33
34
35
36
37
try:
    from outlines.fsm.json_schema import convert_json_schema_to_str
except ImportError:
    # Before outlines 0.0.47, convert_json_schema_to_str is under
    # outlines.integrations.utils
    from outlines.integrations.utils import convert_json_schema_to_str

38
39
40
41
42
43
44
from sglang.srt.conversation import (
    Conversation,
    SeparatorStyle,
    chat_template_exists,
    generate_chat_conv,
    register_conv_template,
)
Ying Sheng's avatar
Ying Sheng committed
45
from sglang.srt.managers.io_struct import EmbeddingReqInput, GenerateReqInput
Mingyi's avatar
Mingyi committed
46
from sglang.srt.openai_api.protocol import (
47
48
    BatchRequest,
    BatchResponse,
49
50
51
52
53
    ChatCompletionRequest,
    ChatCompletionResponse,
    ChatCompletionResponseChoice,
    ChatCompletionResponseStreamChoice,
    ChatCompletionStreamResponse,
54
    ChatCompletionTokenLogprob,
55
    ChatMessage,
56
    ChoiceLogprobs,
57
58
59
60
61
62
    CompletionRequest,
    CompletionResponse,
    CompletionResponseChoice,
    CompletionResponseStreamChoice,
    CompletionStreamResponse,
    DeltaMessage,
Ying Sheng's avatar
Ying Sheng committed
63
    EmbeddingObject,
64
65
    EmbeddingRequest,
    EmbeddingResponse,
66
    ErrorResponse,
67
    FileDeleteResponse,
68
69
    FileRequest,
    FileResponse,
70
    LogProbs,
71
    TopLogprob,
72
73
74
    UsageInfo,
)

75
76
logger = logging.getLogger(__name__)

77
78
chat_template_name = None

Liangsheng Yin's avatar
Liangsheng Yin committed
79

80
81
82
83
84
85
86
87
88
89
class FileMetadata:
    def __init__(self, filename: str, purpose: str):
        self.filename = filename
        self.purpose = purpose


# In-memory storage for batch jobs and files
batch_storage: Dict[str, BatchResponse] = {}
file_id_request: Dict[str, FileMetadata] = {}
file_id_response: Dict[str, FileResponse] = {}
90
# map file id to file path in SGLang backend
91
92
93
94
95
96
97
file_id_storage: Dict[str, str] = {}


# backend storage directory
storage_dir = None


98
99
100
def create_error_response(
    message: str,
    err_type: str = "BadRequestError",
101
102
103
104
    status_code: HTTPStatus = HTTPStatus.BAD_REQUEST,
):
    error = ErrorResponse(message=message, type=err_type, code=status_code.value)
    return JSONResponse(content=error.model_dump(), status_code=error.code)
105
106
107
108
109


def create_streaming_error_response(
    message: str,
    err_type: str = "BadRequestError",
110
111
112
    status_code: HTTPStatus = HTTPStatus.BAD_REQUEST,
) -> str:
    error = ErrorResponse(message=message, type=err_type, code=status_code.value)
113
114
115
116
    json_str = json.dumps({"error": error.model_dump()})
    return json_str


117
def load_chat_template_for_openai_api(tokenizer_manager, chat_template_arg):
118
119
    global chat_template_name

120
121
122
    logger.info(
        f"Use chat template for the OpenAI-compatible API server: {chat_template_arg}"
    )
123
124
125
126
127
128
    if not chat_template_exists(chat_template_arg):
        if not os.path.exists(chat_template_arg):
            raise RuntimeError(
                f"Chat template {chat_template_arg} is not a built-in template name "
                "or a valid chat template file path."
            )
129
130
131
132
133
        if chat_template_arg.endswith(".jinja"):
            with open(chat_template_arg, "r") as f:
                chat_template = "".join(f.readlines()).strip("\n")
            tokenizer_manager.tokenizer.chat_template = chat_template.replace(
                "\\n", "\n"
134
            )
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
            chat_template_name = None
        else:
            assert chat_template_arg.endswith(
                ".json"
            ), "unrecognized format of chat template file"
            with open(chat_template_arg, "r") as filep:
                template = json.load(filep)
                try:
                    sep_style = SeparatorStyle[template["sep_style"]]
                except KeyError:
                    raise ValueError(
                        f"Unknown separator style: {template['sep_style']}"
                    ) from None
                register_conv_template(
                    Conversation(
                        name=template["name"],
                        system_template=template["system"] + "\n{system_message}",
                        system_message=template.get("system_message", ""),
                        roles=(template["user"], template["assistant"]),
                        sep_style=sep_style,
                        sep=template.get("sep", "\n"),
                        stop_str=template["stop_str"],
                    ),
                    override=True,
                )
            chat_template_name = template["name"]
161
162
163
164
    else:
        chat_template_name = chat_template_arg


165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
async def v1_files_create(file: UploadFile, purpose: str, file_storage_pth: str = None):
    try:
        global storage_dir
        if file_storage_pth:
            storage_dir = file_storage_pth
        # Read the file content
        file_content = await file.read()

        # Create an instance of RequestBody
        request_body = FileRequest(file=file_content, purpose=purpose)

        # Save the file to the sglang_oai_storage directory
        os.makedirs(storage_dir, exist_ok=True)
        file_id = f"backend_input_file-{uuid.uuid4()}"
        filename = f"{file_id}.jsonl"
        file_path = os.path.join(storage_dir, filename)

        with open(file_path, "wb") as f:
            f.write(request_body.file)

        # add info to global file map
        file_id_request[file_id] = FileMetadata(filename=file.filename, purpose=purpose)
        file_id_storage[file_id] = file_path

        # Return the response in the required format
        response = FileResponse(
            id=file_id,
            bytes=len(request_body.file),
            created_at=int(time.time()),
            filename=file.filename,
            purpose=request_body.purpose,
        )
        file_id_response[file_id] = response

        return response
    except ValidationError as e:
        return {"error": "Invalid input", "details": e.errors()}


204
205
206
207
208
209
210
211
212
213
214
215
216
217
async def v1_delete_file(file_id: str):
    # Retrieve the file job from the in-memory storage
    file_response = file_id_response.get(file_id)
    if file_response is None:
        raise HTTPException(status_code=404, detail="File not found")
    file_path = file_id_storage.get(file_id)
    if file_path is None:
        raise HTTPException(status_code=404, detail="File not found")
    os.remove(file_path)
    del file_id_response[file_id]
    del file_id_storage[file_id]
    return FileDeleteResponse(id=file_id, deleted=True)


218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
async def v1_batches(tokenizer_manager, raw_request: Request):
    try:
        body = await raw_request.json()

        batch_request = BatchRequest(**body)

        batch_id = f"batch_{uuid.uuid4()}"

        # Create an instance of BatchResponse
        batch_response = BatchResponse(
            id=batch_id,
            endpoint=batch_request.endpoint,
            input_file_id=batch_request.input_file_id,
            completion_window=batch_request.completion_window,
            created_at=int(time.time()),
            metadata=batch_request.metadata,
        )

        batch_storage[batch_id] = batch_response

        # Start processing the batch asynchronously
        asyncio.create_task(process_batch(tokenizer_manager, batch_id, batch_request))

        # Return the initial batch_response
        return batch_response

    except ValidationError as e:
        return {"error": "Invalid input", "details": e.errors()}
    except Exception as e:
        return {"error": str(e)}


async def process_batch(tokenizer_manager, batch_id: str, batch_request: BatchRequest):
    try:
        # Update the batch status to "in_progress"
        batch_storage[batch_id].status = "in_progress"
        batch_storage[batch_id].in_progress_at = int(time.time())

        # Retrieve the input file content
        input_file_request = file_id_request.get(batch_request.input_file_id)
        if not input_file_request:
            raise ValueError("Input file not found")

        # Parse the JSONL file and process each request
        input_file_path = file_id_storage.get(batch_request.input_file_id)
        with open(input_file_path, "r", encoding="utf-8") as f:
            lines = f.readlines()

        total_requests = len(lines)
        completed_requests = 0
        failed_requests = 0

        all_ret = []
        end_point = batch_storage[batch_id].endpoint
        file_request_list = []
        all_requests = []
274
        request_ids = []
275
276
277
278
        for line in lines:
            request_data = json.loads(line)
            file_request_list.append(request_data)
            body = request_data["body"]
279
            request_ids.append(request_data["custom_id"])
280
281
282
283
284
285

            # Although streaming is supported for standalone completions, it is not supported in
            # batch mode (multiple completions in single request).
            if body.get("stream", False):
                raise ValueError("Streaming requests are not supported in batch mode")

286
287
288
289
            if end_point == "/v1/chat/completions":
                all_requests.append(ChatCompletionRequest(**body))
            elif end_point == "/v1/completions":
                all_requests.append(CompletionRequest(**body))
290

291
292
        if end_point == "/v1/chat/completions":
            adapted_request, request = v1_chat_generate_request(
293
                all_requests, tokenizer_manager, request_ids=request_ids
294
295
            )
        elif end_point == "/v1/completions":
296
297
298
299
            adapted_request, request = v1_generate_request(
                all_requests, request_ids=request_ids
            )

300
301
302
303
304
305
306
        try:
            ret = await tokenizer_manager.generate_request(adapted_request).__anext__()
            if not isinstance(ret, list):
                ret = [ret]
            if end_point == "/v1/chat/completions":
                responses = v1_chat_generate_response(request, ret, to_file=True)
            else:
yichuan~'s avatar
yichuan~ committed
307
308
309
                responses = v1_generate_response(
                    request, ret, tokenizer_manager, to_file=True
                )
310
311
312
313
314
315
316
317
318
319
320
321

        except Exception as e:
            error_json = {
                "id": f"batch_req_{uuid.uuid4()}",
                "custom_id": request_data.get("custom_id"),
                "response": None,
                "error": {"message": str(e)},
            }
            all_ret.append(error_json)
            failed_requests += len(file_request_list)

        for idx, response in enumerate(responses):
322
            # the batch_req here can be changed to be named within a batch granularity
323
324
325
326
327
328
329
330
            response_json = {
                "id": f"batch_req_{uuid.uuid4()}",
                "custom_id": file_request_list[idx].get("custom_id"),
                "response": response,
                "error": None,
            }
            all_ret.append(response_json)
            completed_requests += 1
331

332
333
334
335
336
337
338
339
340
341
342
343
        # Write results to a new file
        output_file_id = f"backend_result_file-{uuid.uuid4()}"
        global storage_dir
        output_file_path = os.path.join(storage_dir, f"{output_file_id}.jsonl")
        with open(output_file_path, "w", encoding="utf-8") as f:
            for ret in all_ret:
                f.write(json.dumps(ret) + "\n")

        # Update batch response with output file information
        retrieve_batch = batch_storage[batch_id]
        retrieve_batch.output_file_id = output_file_id
        file_id_storage[output_file_id] = output_file_path
344
345
346
347
348
349
350
        file_id_response[output_file_id] = FileResponse(
            id=output_file_id,
            bytes=os.path.getsize(output_file_path),
            created_at=int(time.time()),
            filename=f"{output_file_id}.jsonl",
            purpose="batch_result",
        )
351
352
353
354
355
356
357
358
359
360
        # Update batch status to "completed"
        retrieve_batch.status = "completed"
        retrieve_batch.completed_at = int(time.time())
        retrieve_batch.request_counts = {
            "total": total_requests,
            "completed": completed_requests,
            "failed": failed_requests,
        }

    except Exception as e:
361
        logger.error("error in SGLang:", e)
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
        # Update batch status to "failed"
        retrieve_batch = batch_storage[batch_id]
        retrieve_batch.status = "failed"
        retrieve_batch.failed_at = int(time.time())
        retrieve_batch.errors = {"message": str(e)}


async def v1_retrieve_batch(batch_id: str):
    # Retrieve the batch job from the in-memory storage
    batch_response = batch_storage.get(batch_id)
    if batch_response is None:
        raise HTTPException(status_code=404, detail="Batch not found")

    return batch_response


378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
async def v1_cancel_batch(tokenizer_manager, batch_id: str):
    # Retrieve the batch job from the in-memory storage
    batch_response = batch_storage.get(batch_id)
    if batch_response is None:
        raise HTTPException(status_code=404, detail="Batch not found")

    # Only do cancal when status is "validating" or "in_progress"
    if batch_response.status in ["validating", "in_progress"]:
        # Start cancelling the batch asynchronously
        asyncio.create_task(
            cancel_batch(
                tokenizer_manager=tokenizer_manager,
                batch_id=batch_id,
                input_file_id=batch_response.input_file_id,
            )
        )

        # Update batch status to "cancelling"
        batch_response.status = "cancelling"

        return batch_response
    else:
        raise HTTPException(
            status_code=500,
            detail=f"Current status is {batch_response.status}, no need to cancel",
        )


async def cancel_batch(tokenizer_manager, batch_id: str, input_file_id: str):
    try:
        # Update the batch status to "cancelling"
        batch_storage[batch_id].status = "cancelling"

        # Retrieve the input file content
        input_file_request = file_id_request.get(input_file_id)
        if not input_file_request:
            raise ValueError("Input file not found")

        # Parse the JSONL file and process each request
        input_file_path = file_id_storage.get(input_file_id)
        with open(input_file_path, "r", encoding="utf-8") as f:
            lines = f.readlines()

        file_request_list = []
        request_ids = []
        for line in lines:
            request_data = json.loads(line)
            file_request_list.append(request_data)
            request_ids.append(request_data["custom_id"])

        # Cancel requests by request_ids
        for rid in request_ids:
            tokenizer_manager.abort_request(rid=rid)

        retrieve_batch = batch_storage[batch_id]
        retrieve_batch.status = "cancelled"

    except Exception as e:
        logger.error("error in SGLang:", e)
        # Update batch status to "failed"
        retrieve_batch = batch_storage[batch_id]
        retrieve_batch.status = "failed"
        retrieve_batch.failed_at = int(time.time())
        retrieve_batch.errors = {"message": str(e)}


444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
async def v1_retrieve_file(file_id: str):
    # Retrieve the batch job from the in-memory storage
    file_response = file_id_response.get(file_id)
    if file_response is None:
        raise HTTPException(status_code=404, detail="File not found")
    return file_response


async def v1_retrieve_file_content(file_id: str):
    file_pth = file_id_storage.get(file_id)
    if not file_pth or not os.path.exists(file_pth):
        raise HTTPException(status_code=404, detail="File not found")

    def iter_file():
        with open(file_pth, mode="rb") as file_like:
            yield from file_like

    return StreamingResponse(iter_file(), media_type="application/octet-stream")


464
465
466
def v1_generate_request(
    all_requests: List[CompletionRequest], request_ids: List[str] = None
):
467
468
    prompts = []
    sampling_params_list = []
469
    return_logprobs = []
470
    logprob_start_lens = []
471
    top_logprobs_nums = []
yichuan~'s avatar
yichuan~ committed
472

473
474
    # NOTE: with openai API, the prompt's logprobs are always not computed
    first_prompt_type = type(all_requests[0].prompt)
475
476
    for request in all_requests:
        assert (
477
            type(request.prompt) is first_prompt_type
478
        ), "All prompts must be of the same type in file input settings"
479
480
481
482
483
        if len(all_requests) > 1 and request.n > 1:
            raise ValueError(
                "Parallel sampling is not supported for completions from files"
            )
        if request.echo and request.logprobs:
484
            logger.warning(
485
486
487
488
489
490
                "Echo is not compatible with logprobs. "
                "To compute logprobs of input prompt, please use SGLang /request API."
            )

    for request in all_requests:
        prompts.append(request.prompt)
491
        return_logprobs.append(request.logprobs is not None and request.logprobs > 0)
492
        logprob_start_lens.append(-1)
493
494
495
        top_logprobs_nums.append(
            request.logprobs if request.logprobs is not None else 0
        )
496
497
498
499
        sampling_params_list.append(
            {
                "temperature": request.temperature,
                "max_new_tokens": request.max_tokens,
500
                "min_new_tokens": request.min_tokens,
501
                "stop": request.stop,
502
                "stop_token_ids": request.stop_token_ids,
503
504
505
                "top_p": request.top_p,
                "presence_penalty": request.presence_penalty,
                "frequency_penalty": request.frequency_penalty,
506
                "repetition_penalty": request.repetition_penalty,
507
                "regex": request.regex,
508
                "json_schema": request.json_schema,
509
510
511
512
513
514
515
516
                "n": request.n,
                "ignore_eos": request.ignore_eos,
            }
        )

    if len(all_requests) == 1:
        prompt = prompts[0]
        sampling_params_list = sampling_params_list[0]
517
        logprob_start_lens = logprob_start_lens[0]
518
519
        return_logprobs = return_logprobs[0]
        top_logprobs_nums = top_logprobs_nums[0]
yichuan~'s avatar
yichuan~ committed
520
        if isinstance(prompt, str) or isinstance(prompt[0], str):
521
522
523
            prompt_kwargs = {"text": prompt}
        else:
            prompt_kwargs = {"input_ids": prompt}
524
    else:
525
        if isinstance(prompts[0], str):
526
527
528
            prompt_kwargs = {"text": prompts}
        else:
            prompt_kwargs = {"input_ids": prompts}
yichuan~'s avatar
yichuan~ committed
529

530
    adapted_request = GenerateReqInput(
531
        **prompt_kwargs,
532
        sampling_params=sampling_params_list,
533
534
        return_logprob=return_logprobs,
        top_logprobs_num=top_logprobs_nums,
535
        logprob_start_len=logprob_start_lens,
536
        return_text_in_logprobs=True,
537
        stream=all_requests[0].stream,
538
        rid=request_ids,
539
    )
yichuan~'s avatar
yichuan~ committed
540

541
542
543
544
545
    if len(all_requests) == 1:
        return adapted_request, all_requests[0]
    return adapted_request, all_requests


yichuan~'s avatar
yichuan~ committed
546
def v1_generate_response(request, ret, tokenizer_manager, to_file=False):
547
548
549
    choices = []
    echo = False

yichuan~'s avatar
yichuan~ committed
550
    if (not isinstance(request, list)) and request.echo:
551
        # TODO: handle the case propmt is token ids
yichuan~'s avatar
yichuan~ committed
552
553
        if isinstance(request.prompt, list) and isinstance(request.prompt[0], str):
            # for the case of multiple str prompts
554
            prompts = request.prompt
yichuan~'s avatar
yichuan~ committed
555
556
557
558
559
560
561
562
563
564
565
566
567
        elif isinstance(request.prompt, list) and isinstance(request.prompt[0], list):
            # for the case of multiple token ids prompts
            prompts = [
                tokenizer_manager.tokenizer.decode(prompt, skip_special_tokens=True)
                for prompt in request.prompt
            ]
        elif isinstance(request.prompt, list) and isinstance(request.prompt[0], int):
            # for the case of single token ids prompt
            prompts = [
                tokenizer_manager.tokenizer.decode(
                    request.prompt, skip_special_tokens=True
                )
            ]
568
        else:
yichuan~'s avatar
yichuan~ committed
569
            # for the case of single str prompt
570
571
572
573
574
            prompts = [request.prompt]
        echo = True

    for idx, ret_item in enumerate(ret):
        text = ret_item["text"]
yichuan~'s avatar
yichuan~ committed
575
        if isinstance(request, list) and request[idx].echo:
576
577
            echo = True
            text = request[idx].prompt + text
yichuan~'s avatar
yichuan~ committed
578
579
580
        if (not isinstance(request, list)) and echo:
            prompt_index = idx // request.n
            text = prompts[prompt_index] + text
581
582

        logprobs = False
yichuan~'s avatar
yichuan~ committed
583
        if isinstance(request, list) and request[idx].logprobs:
584
            logprobs = True
yichuan~'s avatar
yichuan~ committed
585
        elif (not isinstance(request, list)) and request.logprobs:
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
            logprobs = True
        if logprobs:
            if echo:
                input_token_logprobs = ret_item["meta_info"]["input_token_logprobs"]
                input_top_logprobs = ret_item["meta_info"]["input_top_logprobs"]
            else:
                input_token_logprobs = None
                input_top_logprobs = None

            logprobs = to_openai_style_logprobs(
                input_token_logprobs=input_token_logprobs,
                input_top_logprobs=input_top_logprobs,
                output_token_logprobs=ret_item["meta_info"]["output_token_logprobs"],
                output_top_logprobs=ret_item["meta_info"]["output_top_logprobs"],
            )
        else:
            logprobs = None

        if to_file:
605
            # to make the choise data json serializable
606
607
608
609
            choice_data = {
                "index": 0,
                "text": text,
                "logprobs": logprobs,
610
611
612
613
                "finish_reason": (
                    ret_item["meta_info"]["finish_reason"]["type"]
                    if ret_item["meta_info"]["finish_reason"]
                    else ""
614
                ),
615
616
617
618
619
620
            }
        else:
            choice_data = CompletionResponseChoice(
                index=idx,
                text=text,
                logprobs=logprobs,
621
622
623
624
                finish_reason=(
                    ret_item["meta_info"]["finish_reason"]["type"]
                    if ret_item["meta_info"]["finish_reason"]
                    else ""
625
                ),
626
627
628
629
630
631
632
633
634
635
636
            )

        choices.append(choice_data)

    if to_file:
        responses = []
        for i, choice in enumerate(choices):
            response = {
                "status_code": 200,
                "request_id": ret[i]["meta_info"]["id"],
                "body": {
637
                    # remain the same but if needed we can change that
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
                    "id": ret[i]["meta_info"]["id"],
                    "object": "text_completion",
                    "created": int(time.time()),
                    "model": request[i].model,
                    "choices": choice,
                    "usage": {
                        "prompt_tokens": ret[i]["meta_info"]["prompt_tokens"],
                        "completion_tokens": ret[i]["meta_info"]["completion_tokens"],
                        "total_tokens": ret[i]["meta_info"]["prompt_tokens"]
                        + ret[i]["meta_info"]["completion_tokens"],
                    },
                    "system_fingerprint": None,
                },
            }
            responses.append(response)
        return responses
    else:
655
656
657
        prompt_tokens = sum(
            ret[i]["meta_info"]["prompt_tokens"] for i in range(0, len(ret), request.n)
        )
658
659
660
661
662
663
        completion_tokens = sum(item["meta_info"]["completion_tokens"] for item in ret)
        response = CompletionResponse(
            id=ret[0]["meta_info"]["id"],
            model=request.model,
            choices=choices,
            usage=UsageInfo(
yichuan~'s avatar
yichuan~ committed
664
                prompt_tokens=prompt_tokens,
665
                completion_tokens=completion_tokens,
yichuan~'s avatar
yichuan~ committed
666
                total_tokens=prompt_tokens + completion_tokens,
667
668
669
670
671
672
673
674
675
            ),
        )
    return response


async def v1_completions(tokenizer_manager, raw_request: Request):
    request_json = await raw_request.json()
    all_requests = [CompletionRequest(**request_json)]
    adapted_request, request = v1_generate_request(all_requests)
676
677
678
679

    if adapted_request.stream:

        async def generate_stream_resp():
680
681
682
683
            stream_buffers = {}
            n_prev_tokens = {}
            prompt_tokens = {}
            completion_tokens = {}
684
685
            try:
                async for content in tokenizer_manager.generate_request(
686
687
                    adapted_request, raw_request
                ):
688
689
690
691
692
                    index = content["index"]

                    stream_buffer = stream_buffers.get(index, "")
                    n_prev_token = n_prev_tokens.get(index, 0)

693
                    text = content["text"]
694
695
                    prompt_tokens[index] = content["meta_info"]["prompt_tokens"]
                    completion_tokens[index] = content["meta_info"]["completion_tokens"]
696
697
698

                    if not stream_buffer:  # The first chunk
                        if request.echo:
yichuan~'s avatar
yichuan~ committed
699
700
701
                            if isinstance(request.prompt, str):
                                # for the case of single str prompts
                                prompts = request.prompt
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
                            elif isinstance(request.prompt, list):
                                if isinstance(request.prompt[0], str):
                                    # for the case of multiple str prompts
                                    prompts = request.prompt[index // request.n]
                                elif isinstance(request.prompt[0], int):
                                    # for the case of single token ids prompt
                                    prompts = tokenizer_manager.tokenizer.decode(
                                        request.prompt, skip_special_tokens=True
                                    )
                                elif isinstance(request.prompt[0], list) and isinstance(
                                    request.prompt[0][0], int
                                ):
                                    # for the case of multiple token ids prompts
                                    prompts = tokenizer_manager.tokenizer.decode(
                                        request.prompt[index // request.n],
                                        skip_special_tokens=True,
                                    )
yichuan~'s avatar
yichuan~ committed
719

720
                            # Prepend prompt in response text.
yichuan~'s avatar
yichuan~ committed
721
                            text = prompts + text
722
723
724
725

                    if request.logprobs:
                        # The first chunk and echo is enabled.
                        if not stream_buffer and request.echo:
726
727
                            input_token_logprobs = content["meta_info"][
                                "input_token_logprobs"
728
                            ]
729
730
                            input_top_logprobs = content["meta_info"][
                                "input_top_logprobs"
731
732
                            ]
                        else:
733
734
                            input_token_logprobs = None
                            input_top_logprobs = None
735
736

                        logprobs = to_openai_style_logprobs(
737
738
739
740
                            input_token_logprobs=input_token_logprobs,
                            input_top_logprobs=input_top_logprobs,
                            output_token_logprobs=content["meta_info"][
                                "output_token_logprobs"
741
                            ][n_prev_token:],
742
743
                            output_top_logprobs=content["meta_info"][
                                "output_top_logprobs"
744
                            ][n_prev_token:],
745
                        )
746
                        n_prev_token = len(
747
                            content["meta_info"]["output_token_logprobs"]
748
                        )
749
                    else:
750
                        logprobs = None
751

752
                    delta = text[len(stream_buffer) :]
Liangsheng Yin's avatar
Liangsheng Yin committed
753
                    stream_buffer = stream_buffer + delta
754
                    choice_data = CompletionResponseStreamChoice(
755
                        index=index,
756
757
                        text=delta,
                        logprobs=logprobs,
758
759
760
761
                        finish_reason=(
                            content["meta_info"]["finish_reason"]["type"]
                            if content["meta_info"]["finish_reason"]
                            else ""
762
                        ),
763
764
765
766
767
768
769
                    )
                    chunk = CompletionStreamResponse(
                        id=content["meta_info"]["id"],
                        object="text_completion",
                        choices=[choice_data],
                        model=request.model,
                    )
770
771
772
773

                    stream_buffers[index] = stream_buffer
                    n_prev_tokens[index] = n_prev_token

774
                    yield f"data: {chunk.model_dump_json()}\n\n"
775
                if request.stream_options and request.stream_options.include_usage:
776
777
778
779
780
781
782
783
                    total_prompt_tokens = sum(
                        tokens
                        for i, tokens in prompt_tokens.items()
                        if i % request.n == 0
                    )
                    total_completion_tokens = sum(
                        tokens for tokens in completion_tokens.values()
                    )
784
                    usage = UsageInfo(
785
786
787
                        prompt_tokens=total_prompt_tokens,
                        completion_tokens=total_completion_tokens,
                        total_tokens=total_prompt_tokens + total_completion_tokens,
788
789
790
791
792
793
794
795
796
797
798
799
                    )

                    final_usage_chunk = CompletionStreamResponse(
                        id=str(uuid.uuid4().hex),
                        choices=[],
                        model=request.model,
                        usage=usage,
                    )
                    final_usage_data = final_usage_chunk.model_dump_json(
                        exclude_unset=True, exclude_none=True
                    )
                    yield f"data: {final_usage_data}\n\n"
800
801
802
            except ValueError as e:
                error = create_streaming_error_response(str(e))
                yield f"data: {error}\n\n"
803
804
            yield "data: [DONE]\n\n"

805
806
807
808
809
        return StreamingResponse(
            generate_stream_resp(),
            media_type="text/event-stream",
            background=tokenizer_manager.create_abort_task(adapted_request),
        )
810
811

    # Non-streaming response.
812
813
    try:
        ret = await tokenizer_manager.generate_request(
814
815
            adapted_request, raw_request
        ).__anext__()
816
817
    except ValueError as e:
        return create_error_response(str(e))
818

819
820
821
    if not isinstance(ret, list):
        ret = [ret]

yichuan~'s avatar
yichuan~ committed
822
    response = v1_generate_response(request, ret, tokenizer_manager)
823
    return response
824

825

826
def v1_chat_generate_request(
827
828
829
    all_requests: List[ChatCompletionRequest],
    tokenizer_manager,
    request_ids: List[str] = None,
830
):
831
    input_ids = []
832
833
    sampling_params_list = []
    image_data_list = []
834
    return_logprobs = []
835
    logprob_start_lens = []
836
    top_logprobs_nums = []
837
    modalities_list = []
838
839
840

    # NOTE: with openai API, the prompt's logprobs are always not computed

841
842
843
844
845
846
847
848
849
    for request in all_requests:
        # Prep the data needed for the underlying GenerateReqInput:
        #  - prompt: The full prompt string.
        #  - stop: Custom stop tokens.
        #  - image_data: None or a list of image strings (URLs or base64 strings).
        #    None skips any image processing in GenerateReqInput.
        if not isinstance(request.messages, str):
            # Apply chat template and its stop strings.
            if chat_template_name is None:
850
851
852
853
854
855
856
857
858
859
860
861
862
                openai_compatible_messages = []
                for message in request.messages:
                    if isinstance(message.content, str):
                        openai_compatible_messages.append(
                            {"role": message.role, "content": message.content}
                        )
                    else:
                        content_list = message.dict()["content"]
                        for content in content_list:
                            if content["type"] == "text":
                                openai_compatible_messages.append(
                                    {"role": message.role, "content": content["text"]}
                                )
863
864
865
866
867
                if openai_compatible_messages[-1]["role"] == "assistant":
                    assistant_prefix = openai_compatible_messages[-1]["content"]
                    openai_compatible_messages = openai_compatible_messages[:-1]
                else:
                    assistant_prefix = None
868
                prompt_ids = tokenizer_manager.tokenizer.apply_chat_template(
869
870
871
                    openai_compatible_messages,
                    tokenize=True,
                    add_generation_prompt=True,
872
                )
873
874
                if assistant_prefix:
                    prompt_ids += tokenizer_manager.tokenizer.encode(assistant_prefix)
875
876
                stop = request.stop
                image_data = None
877
                modalities = []
878
            else:
879
880
881
                conv = generate_chat_conv(request, chat_template_name)
                prompt = conv.get_prompt()
                image_data = conv.image_data
882
                modalities = conv.modalities
883
884
885
886
887
888
                stop = conv.stop_str or []
                if request.stop:
                    if isinstance(request.stop, str):
                        stop.append(request.stop)
                    else:
                        stop.extend(request.stop)
889
                prompt_ids = tokenizer_manager.tokenizer.encode(prompt)
890
        else:
891
            # Use the raw prompt and stop strings if the messages is already a string.
yichuan~'s avatar
yichuan~ committed
892
            prompt_ids = request.messages
893
894
            stop = request.stop
            image_data = None
895
            modalities = []
896
        input_ids.append(prompt_ids)
897
        return_logprobs.append(request.logprobs)
898
        logprob_start_lens.append(-1)
899
        top_logprobs_nums.append(request.top_logprobs or 0)
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919

        sampling_params = {
            "temperature": request.temperature,
            "max_new_tokens": request.max_tokens,
            "min_new_tokens": request.min_tokens,
            "stop": stop,
            "stop_token_ids": request.stop_token_ids,
            "top_p": request.top_p,
            "presence_penalty": request.presence_penalty,
            "frequency_penalty": request.frequency_penalty,
            "repetition_penalty": request.repetition_penalty,
            "regex": request.regex,
            "n": request.n,
        }
        if request.response_format and request.response_format.type == "json_schema":
            sampling_params["json_schema"] = convert_json_schema_to_str(
                request.response_format.json_schema.schema_
            )
        sampling_params_list.append(sampling_params)

920
        image_data_list.append(image_data)
921
        modalities_list.extend(modalities)
922
    if len(all_requests) == 1:
923
        input_ids = input_ids[0]
yichuan~'s avatar
yichuan~ committed
924
925
926
927
        if isinstance(input_ids, str):
            prompt_kwargs = {"text": input_ids}
        else:
            prompt_kwargs = {"input_ids": input_ids}
928
        sampling_params_list = sampling_params_list[0]
929
        image_data_list = image_data_list[0]
930
        return_logprobs = return_logprobs[0]
931
        logprob_start_lens = logprob_start_lens[0]
932
        top_logprobs_nums = top_logprobs_nums[0]
933
        modalities_list = modalities_list[:1]
yichuan~'s avatar
yichuan~ committed
934
935
936
937
938
    else:
        if isinstance(input_ids[0], str):
            prompt_kwargs = {"text": input_ids}
        else:
            prompt_kwargs = {"input_ids": input_ids}
939

940
    adapted_request = GenerateReqInput(
yichuan~'s avatar
yichuan~ committed
941
        **prompt_kwargs,
942
        image_data=image_data_list,
943
        sampling_params=sampling_params_list,
944
        return_logprob=return_logprobs,
945
        logprob_start_len=logprob_start_lens,
946
947
948
        top_logprobs_num=top_logprobs_nums,
        stream=all_requests[0].stream,
        return_text_in_logprobs=True,
949
        rid=request_ids,
950
        modalities=modalities_list,
951
    )
952
953
954
    if len(all_requests) == 1:
        return adapted_request, all_requests[0]
    return adapted_request, all_requests
955

956

957
958
959
960
def v1_chat_generate_response(request, ret, to_file=False):
    choices = []

    for idx, ret_item in enumerate(ret):
961
        logprobs = False
yichuan~'s avatar
yichuan~ committed
962
        if isinstance(request, list) and request[idx].logprobs:
963
            logprobs = True
yichuan~'s avatar
yichuan~ committed
964
        elif (not isinstance(request, list)) and request.logprobs:
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
            logprobs = True
        if logprobs:
            logprobs = to_openai_style_logprobs(
                output_token_logprobs=ret_item["meta_info"]["output_token_logprobs"],
                output_top_logprobs=ret_item["meta_info"]["output_top_logprobs"],
            )
            token_logprobs = []
            for token, logprob in zip(logprobs.tokens, logprobs.token_logprobs):
                token_bytes = list(token.encode("utf-8"))
                top_logprobs = []
                if logprobs.top_logprobs:
                    for top_token, top_logprob in logprobs.top_logprobs[0].items():
                        top_token_bytes = list(top_token.encode("utf-8"))
                        top_logprobs.append(
                            TopLogprob(
                                token=top_token,
                                bytes=top_token_bytes,
                                logprob=top_logprob,
                            )
                        )
                token_logprobs.append(
                    ChatCompletionTokenLogprob(
                        token=token,
                        bytes=token_bytes,
                        logprob=logprob,
                        top_logprobs=top_logprobs,
                    )
                )

            choice_logprobs = ChoiceLogprobs(content=token_logprobs)
        else:
            choice_logprobs = None
997

998
        if to_file:
999
            # to make the choice data json serializable
1000
1001
1002
            choice_data = {
                "index": 0,
                "message": {"role": "assistant", "content": ret_item["text"]},
1003
                "logprobs": choice_logprobs,
1004
1005
1006
1007
                "finish_reason": (
                    ret_item["meta_info"]["finish_reason"]["type"]
                    if ret_item["meta_info"]["finish_reason"]
                    else ""
1008
                ),
1009
            }
1010
        else:
1011
1012
1013
            choice_data = ChatCompletionResponseChoice(
                index=idx,
                message=ChatMessage(role="assistant", content=ret_item["text"]),
1014
                logprobs=choice_logprobs,
1015
1016
1017
1018
                finish_reason=(
                    ret_item["meta_info"]["finish_reason"]["type"]
                    if ret_item["meta_info"]["finish_reason"]
                    else ""
1019
                ),
1020
1021
1022
            )

        choices.append(choice_data)
1023

1024
1025
1026
1027
1028
1029
1030
1031
    if to_file:
        responses = []

        for i, choice in enumerate(choices):
            response = {
                "status_code": 200,
                "request_id": ret[i]["meta_info"]["id"],
                "body": {
1032
                    # remain the same but if needed we can change that
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
                    "id": ret[i]["meta_info"]["id"],
                    "object": "chat.completion",
                    "created": int(time.time()),
                    "model": request[i].model,
                    "choices": choice,
                    "usage": {
                        "prompt_tokens": ret[i]["meta_info"]["prompt_tokens"],
                        "completion_tokens": ret[i]["meta_info"]["completion_tokens"],
                        "total_tokens": ret[i]["meta_info"]["prompt_tokens"]
                        + ret[i]["meta_info"]["completion_tokens"],
                    },
                    "system_fingerprint": None,
                },
            }
            responses.append(response)
        return responses
1049
    else:
1050
1051
1052
1053
        prompt_tokens = sum(
            ret[i]["meta_info"]["prompt_tokens"] for i in range(0, len(ret), request.n)
        )
        completion_tokens = sum(item["meta_info"]["completion_tokens"] for item in ret)
1054
1055
1056
1057
1058
        response = ChatCompletionResponse(
            id=ret[0]["meta_info"]["id"],
            model=request.model,
            choices=choices,
            usage=UsageInfo(
1059
1060
1061
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                total_tokens=prompt_tokens + completion_tokens,
1062
1063
1064
            ),
        )
        return response
1065

1066
1067
1068
1069
1070

async def v1_chat_completions(tokenizer_manager, raw_request: Request):
    request_json = await raw_request.json()
    all_requests = [ChatCompletionRequest(**request_json)]
    adapted_request, request = v1_chat_generate_request(all_requests, tokenizer_manager)
1071
1072
1073
1074

    if adapted_request.stream:

        async def generate_stream_resp():
1075
1076
1077
1078
1079
            is_firsts = {}
            stream_buffers = {}
            n_prev_tokens = {}
            prompt_tokens = {}
            completion_tokens = {}
1080
            try:
1081
1082
1083
                async for content in tokenizer_manager.generate_request(
                    adapted_request, raw_request
                ):
1084
1085
1086
1087
1088
1089
1090
1091
                    index = content["index"]

                    is_first = is_firsts.get(index, True)
                    stream_buffer = stream_buffers.get(index, "")
                    n_prev_token = n_prev_tokens.get(index, 0)

                    prompt_tokens[index] = content["meta_info"]["prompt_tokens"]
                    completion_tokens[index] = content["meta_info"]["completion_tokens"]
yichuan~'s avatar
yichuan~ committed
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
                    if request.logprobs:
                        logprobs = to_openai_style_logprobs(
                            output_token_logprobs=content["meta_info"][
                                "output_token_logprobs"
                            ][n_prev_token:],
                            output_top_logprobs=content["meta_info"][
                                "output_top_logprobs"
                            ][n_prev_token:],
                        )

                        n_prev_token = len(
                            content["meta_info"]["output_token_logprobs"]
                        )
                        token_logprobs = []
                        for token, logprob in zip(
                            logprobs.tokens, logprobs.token_logprobs
                        ):
                            token_bytes = list(token.encode("utf-8"))
                            top_logprobs = []
                            if logprobs.top_logprobs:
                                for top_token, top_logprob in logprobs.top_logprobs[
                                    0
                                ].items():
                                    top_token_bytes = list(top_token.encode("utf-8"))
                                    top_logprobs.append(
                                        TopLogprob(
                                            token=top_token,
                                            bytes=top_token_bytes,
                                            logprob=top_logprob,
                                        )
                                    )
                            token_logprobs.append(
                                ChatCompletionTokenLogprob(
                                    token=token,
                                    bytes=token_bytes,
                                    logprob=logprob,
                                    top_logprobs=top_logprobs,
                                )
                            )

                        choice_logprobs = ChoiceLogprobs(content=token_logprobs)

                    else:
                        choice_logprobs = None

1137
1138
1139
1140
                    if is_first:
                        # First chunk with role
                        is_first = False
                        choice_data = ChatCompletionResponseStreamChoice(
1141
                            index=index,
1142
                            delta=DeltaMessage(role="assistant"),
1143
1144
1145
1146
                            finish_reason=(
                                content["meta_info"]["finish_reason"]["type"]
                                if content["meta_info"]["finish_reason"]
                                else ""
1147
                            ),
yichuan~'s avatar
yichuan~ committed
1148
                            logprobs=choice_logprobs,
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
                        )
                        chunk = ChatCompletionStreamResponse(
                            id=content["meta_info"]["id"],
                            choices=[choice_data],
                            model=request.model,
                        )
                        yield f"data: {chunk.model_dump_json()}\n\n"

                    text = content["text"]
                    delta = text[len(stream_buffer) :]
Liangsheng Yin's avatar
Liangsheng Yin committed
1159
                    stream_buffer = stream_buffer + delta
1160
                    choice_data = ChatCompletionResponseStreamChoice(
1161
                        index=index,
1162
                        delta=DeltaMessage(content=delta),
1163
1164
1165
1166
                        finish_reason=(
                            content["meta_info"]["finish_reason"]["type"]
                            if content["meta_info"]["finish_reason"]
                            else ""
1167
                        ),
yichuan~'s avatar
yichuan~ committed
1168
                        logprobs=choice_logprobs,
1169
1170
1171
1172
1173
1174
                    )
                    chunk = ChatCompletionStreamResponse(
                        id=content["meta_info"]["id"],
                        choices=[choice_data],
                        model=request.model,
                    )
1175
1176
1177
1178
1179

                    is_firsts[index] = is_first
                    stream_buffers[index] = stream_buffer
                    n_prev_tokens[index] = n_prev_token

1180
                    yield f"data: {chunk.model_dump_json()}\n\n"
1181
                if request.stream_options and request.stream_options.include_usage:
1182
1183
1184
1185
1186
1187
1188
1189
                    total_prompt_tokens = sum(
                        tokens
                        for i, tokens in prompt_tokens.items()
                        if i % request.n == 0
                    )
                    total_completion_tokens = sum(
                        tokens for tokens in completion_tokens.values()
                    )
1190
                    usage = UsageInfo(
1191
1192
1193
                        prompt_tokens=total_prompt_tokens,
                        completion_tokens=total_completion_tokens,
                        total_tokens=total_prompt_tokens + total_completion_tokens,
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
                    )

                    final_usage_chunk = ChatCompletionStreamResponse(
                        id=str(uuid.uuid4().hex),
                        choices=[],
                        model=request.model,
                        usage=usage,
                    )
                    final_usage_data = final_usage_chunk.model_dump_json(
                        exclude_unset=True, exclude_none=True
                    )
                    yield f"data: {final_usage_data}\n\n"
1206
1207
1208
            except ValueError as e:
                error = create_streaming_error_response(str(e))
                yield f"data: {error}\n\n"
1209
1210
            yield "data: [DONE]\n\n"

1211
1212
1213
1214
1215
        return StreamingResponse(
            generate_stream_resp(),
            media_type="text/event-stream",
            background=tokenizer_manager.create_abort_task(adapted_request),
        )
1216
1217

    # Non-streaming response.
1218
1219
    try:
        ret = await tokenizer_manager.generate_request(
1220
1221
            adapted_request, raw_request
        ).__anext__()
1222
1223
    except ValueError as e:
        return create_error_response(str(e))
1224
1225
1226
    if not isinstance(ret, list):
        ret = [ret]

1227
    response = v1_chat_generate_response(request, ret)
1228

1229
1230
1231
    return response


1232
1233
1234
def v1_embedding_request(all_requests, tokenizer_manager):
    prompts = []
    sampling_params_list = []
Ying Sheng's avatar
Ying Sheng committed
1235
    first_prompt_type = type(all_requests[0].input)
1236
1237

    for request in all_requests:
Ying Sheng's avatar
Ying Sheng committed
1238
        prompt = request.input
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
        assert (
            type(prompt) == first_prompt_type
        ), "All prompts must be of the same type in file input settings"
        prompts.append(prompt)

    if len(all_requests) == 1:
        prompt = prompts[0]
        if isinstance(prompt, str) or isinstance(prompt[0], str):
            prompt_kwargs = {"text": prompt}
        else:
            prompt_kwargs = {"input_ids": prompt}
    else:
        if isinstance(prompts[0], str) or isinstance(propmt[0][0], str):
            prompt_kwargs = {"text": prompts}
        else:
            prompt_kwargs = {"input_ids": prompts}

    adapted_request = EmbeddingReqInput(
        **prompt_kwargs,
    )

    if len(all_requests) == 1:
        return adapted_request, all_requests[0]
    return adapted_request, all_requests


Ying Sheng's avatar
Ying Sheng committed
1265
1266
1267
def v1_embedding_response(ret, model_path, to_file=False):
    embedding_objects = []
    prompt_tokens = 0
1268
    for idx, ret_item in enumerate(ret):
Ying Sheng's avatar
Ying Sheng committed
1269
1270
1271
        embedding_objects.append(
            EmbeddingObject(
                embedding=ret[idx]["embedding"],
1272
1273
1274
                index=idx,
            )
        )
Ying Sheng's avatar
Ying Sheng committed
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
        prompt_tokens += ret[idx]["meta_info"]["prompt_tokens"]

    return EmbeddingResponse(
        data=embedding_objects,
        model=model_path,
        usage=UsageInfo(
            prompt_tokens=prompt_tokens,
            total_tokens=prompt_tokens,
        ),
    )
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301


async def v1_embeddings(tokenizer_manager, raw_request: Request):
    request_json = await raw_request.json()
    all_requests = [EmbeddingRequest(**request_json)]
    adapted_request, request = v1_embedding_request(all_requests, tokenizer_manager)

    try:
        ret = await tokenizer_manager.generate_request(
            adapted_request, raw_request
        ).__anext__()
    except ValueError as e:
        return create_error_response(str(e))

    if not isinstance(ret, list):
        ret = [ret]

Ying Sheng's avatar
Ying Sheng committed
1302
    response = v1_embedding_response(ret, tokenizer_manager.model_path)
1303
1304
1305
1306

    return response


1307
def to_openai_style_logprobs(
1308
1309
1310
1311
    input_token_logprobs=None,
    output_token_logprobs=None,
    input_top_logprobs=None,
    output_top_logprobs=None,
1312
1313
1314
1315
1316
1317
1318
1319
):
    ret_logprobs = LogProbs()

    def append_token_logprobs(token_logprobs):
        for logprob, _, token_text in token_logprobs:
            ret_logprobs.tokens.append(token_text)
            ret_logprobs.token_logprobs.append(logprob)

1320
            # Not supported yet
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
            ret_logprobs.text_offset.append(-1)

    def append_top_logprobs(top_logprobs):
        for tokens in top_logprobs:
            if tokens is not None:
                ret_logprobs.top_logprobs.append(
                    {token[2]: token[0] for token in tokens}
                )
            else:
                ret_logprobs.top_logprobs.append(None)

1332
1333
1334
1335
1336
1337
1338
1339
    if input_token_logprobs is not None:
        append_token_logprobs(input_token_logprobs)
    if output_token_logprobs is not None:
        append_token_logprobs(output_token_logprobs)
    if input_top_logprobs is not None:
        append_top_logprobs(input_top_logprobs)
    if output_top_logprobs is not None:
        append_top_logprobs(output_top_logprobs)
1340

Liangsheng Yin's avatar
Liangsheng Yin committed
1341
    return ret_logprobs