Repository: LAION-AI/Open-Assistant
File: inference/worker/basic_hf_server.py (Python)

Code before:

@app.middleware("http")
async def log_exceptions(request: fastapi.Request, call_next):
    try:
        response = await call_next(request)
    except Exception:
        logger.exception("Exception in request")
        raise
    return response


def terminate_server(signum, frame):
    logger.warning(f"Signal {signum}. Terminating server...")
    sys.exit(0)


model: transformers.PreTrainedModel
tokenizer: transformers.PreTrainedTokenizer
fully_loaded: bool = False
model_input_queue: Queue = Queue()


def model_thread():
    request: interface.GenerateStreamRequest
    output_queue: Queue
    eos_token = ""
    if hasattr(tokenizer, "eos_token"):
        eos_token = tokenizer.eos_token
    while True:
        request, output_queue = model_input_queue.get()
        try:
            prompt = request.inputs
            params = request.parameters.dict()
            params.pop("seed")
            params.pop("stop")
            params.pop("details")

            def print_text(text: str, *, end="\n", flush=False):
                if eos_token and text.endswith(eos_token):
                    text = text[: -len(eos_token)]
                if not text:
                    return
                stream_response = interface.GenerateStreamResponse(
                    token=interface.Token(text=text),
                )
                output_queue.put_nowait(stream_response)

            streamer = hf_streamer.HFStreamer(tokenizer=tokenizer, printer=print_text)

            with torch.no_grad():
                ids = tokenizer.encode(prompt, return_tensors="pt")

Review comment (suggestion):
Just a preference really but I feel a one-liner is clearer here

`eos_token = tokenizer.eos_token if hasattr(tokenizer, "eos_token") else ""`

Code after:


@app.middleware("http")
async def log_exceptions(request: fastapi.Request, call_next):
    try:
        response = await call_next(request)
    except Exception:
        logger.exception("Exception in request")
        raise
    return response


model_loaded: bool = False
fully_loaded: bool = False
model_input_queue: Queue = Queue()


def model_thread():
    model: transformers.PreTrainedModel
    tokenizer: transformers.PreTrainedTokenizer
    model, tokenizer, decode_token = load_models()

    request: interface.GenerateStreamRequest
    output_queue: Queue
    eos_token_id = tokenizer.eos_token_id if hasattr(tokenizer, "eos_token_id") else None
    while True:
        request, output_queue = model_input_queue.get()
        try:
            prompt = request.inputs
            params = request.parameters.dict()
            seed = params.pop("seed")
            params.pop("stop")
            params.pop("details")

            last_token_id = None  # need to delay by 1 to simulate tgi

            def print_text(token_id: int):
                nonlocal last_token_id
                if last_token_id is not None:
                    text = decode_token(last_token_id)
                    stream_response = interface.GenerateStreamResponse(
                        token=interface.Token(text=text, id=last_token_id),
                    )
                    output_queue.put_nowait(stream_response)
                last_token_id = token_id

            with torch.no_grad():
                ids = tokenizer.encode(prompt, return_tensors="pt", add_special_tokens=False)
                streamer = hf_streamer.HFStreamer(input_ids=ids, printer=print_text)
                ids = ids.to(model.device)
                output = model.generate(ids, **params, streamer=streamer, eos_token_id=eos_token_id)