Function generate_stream

demo/inference.py:185–241

(
        model, tokenizer, image_processor, params, device
)

@torch.inference_mode()
def generate_stream(
        model, tokenizer, image_processor, params, device
):
    prompt = params["prompt"]
    images = params.get("images", None)
    videos = params.get("videos", None)
    temperature = float(params.get("temperature", 0.7))
    max_new_tokens = int(params.get("max_new_tokens", 1024))

    num_queries = model.config.num_query_tokens

    stop_words = ["Human: ", "Assistant: ", "###", "\n\n"]
    stop_words_ids = [tokenizer(stop_word, return_tensors='pt')['input_ids'].squeeze() for stop_word in stop_words]
    stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])

    generation_config = GenerationConfig(
        bos_token_id=1,
        do_sample=True,
        temperature=temperature,
        max_new_tokens=max_new_tokens,
        stopping_criteria=stopping_criteria
    )

    pixel_values = None
    if images is not None:
        pixel_values = load_image(images).to(device)  # only support one image
        image_query = DEFAULT_IMG_START_TOKEN + DEFAULT_IMG_END_TOKEN
        prompt = prompt.replace("<image>", image_query)

    elif videos is not None:
        pixel_values = load_video(videos).to(device)
        video_query = DEFAULT_VIDEO_START_TOKEN + DEFAULT_VIDEO_END_TOKEN
        prompt = prompt.replace("<video>", video_query)

    model_inputs = tokenizer([prompt], return_tensors="pt")
    model_inputs.pop("token_type_ids", None)

    if pixel_values is not None:
        model_inputs["pixel_values"] = pixel_values

        generation_output = model.generate(
            **model_inputs,
            generation_config=generation_config,
            return_dict_in_generate=True,
            output_scores=True
        )
    else:
        generation_output = model.language_model.generate(
            **model_inputs,
            generation_config=generation_config,
            return_dict_in_generate=True,
            output_scores=True
        )

    preds = generation_output.sequences
    outputs = tokenizer.batch_decode(preds, skip_special_tokens=True)
    return outputs