kezouke
commented
4 months ago Server Code:
import argparse
import uvicorn
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, Response
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.engine.async_llm_engine import AsyncLLMEngine
from vllm.sampling_params import SamplingParams
from vllm.utils import random_uuid
from vllm import TokensPrompt
from aiostream import stream
import torch
TIMEOUT_KEEP_ALIVE = 5 # seconds.
TIMEOUT_TO_PREVENT_DEADLOCK = 1 # seconds.
app = FastAPI()
engine = None
async def fast_logprobs(common_part, tokenized_candidates, k, tokenizer):
# Initialize sampling parameters for generating text
sampling_params = SamplingParams(max_tokens=1, prompt_logprobs=1)
# Dictionary to store unique prompts and corresponding tokenized candidates
input_prompts = {}
# Placeholder for results
results = None
# Iterate over each tokenized candidate
for candidate in tokenized_candidates:
# Calculate the length of the candidate excluding padding tokens
candidate_length = torch.sum(candidate!= tokenizer.pad_token_id, dim=-1)
# Combine the common part with the candidate up to its actual length
prompt_token_ids = common_part + candidate.tolist()[:candidate_length]
# Store the combined prompt and the trimmed candidate in the dictionary
input_prompts[tuple(prompt_token_ids)] = candidate.tolist()[:candidate_length]
# Iterate over each unique prompt
for input_prompt in input_prompts.keys():
# Generate a unique identifier for the request
request_id = random_uuid()
# Prepare the prompt for the model
tokens_prompt: TokensPrompt = {
'prompt_token_ids': list(input_prompt),
'multi_modal_data': None
}
# Use the model to generate logits for the prompt
res = engine.generate(request_id=request_id,
inputs=tokens_prompt,
sampling_params=sampling_params)
# Merge the new result with existing results
if not results:
results = res
else:
results = stream.merge(results, res)
# Convert the results into a list asynchronously
results = [item async for item in results]
# Map each candidate to its cumulative log probability
candidate_log_prob_map = {}
# Iterate over each result
for result in results:
# Retrieve the prompt token IDs from the result
prompt_token_ids = result.prompt_token_ids
# Get the corresponding candidate from the input_prompts dictionary
candidate = input_prompts[tuple(prompt_token_ids)]
# Initialize the cumulative log probability
cum_log_prob = 0
# Calculate the cumulative log probability for the candidate
for token_id, log_prob in zip(candidate, result.prompt_logprobs[-len(candidate):]):
cum_log_prob += log_prob[token_id].logprob
# Store the cumulative log probability in the map
candidate_log_prob_map[tuple(candidate)] = cum_log_prob
# Calculate the normalized scores for each candidate
scores = []
for candidate in tokenized_candidates:
# Calculate the length of the candidate excluding padding tokens
candidate_length = torch.sum(candidate!= tokenizer.pad_token_id, dim=-1)
# Retrieve the cumulative log probability for the candidate
cum_log_probability = candidate_log_prob_map[tuple(candidate.tolist()[:candidate_length])]
# Normalize the score by dividing by the candidate length
scores.append(cum_log_probability / candidate_length)
# Convert the scores into a tensor
scores = torch.tensor(scores)
# Find the indices of the top k scores
top_k_indices = torch.topk(scores, k=k, largest=True, sorted=True).indices
# Return the top k candidates based on their scores
return [tokenized_candidates[i] for i in top_k_indices]
@app.post("/generate")
async def generate(request: Request) -> Response:
"""Generate completion for the request.
The request should be a JSON object with the following fields:
- prompt: the prompt to use for the generation.
- schema: the JSON schema to use for the generation (if regex is not provided).
- regex: the regex to use for the generation (if schema is not provided).
- stream: whether to stream the results or not.
- other fields: the sampling parameters (See `SamplingParams` for details).
"""
assert engine is not None
request_dict = await request.json()
prompt = request_dict.pop('prompt',None)
substring_context = request_dict.pop('substring',None)
if substring_context is not None:
tokenizer = engine.engine.get_tokenizer()
if not tokenizer.pad_token:
tokenizer.pad_token = tokenizer.eos_token
tokenized_prompt = tokenizer(prompt,
return_tensors="pt",
padding=True,
add_special_tokens=False)['input_ids']
last_part = tokenizer.decode(tokenized_prompt[0, -2:])
common_part = tokenizer.decode(tokenized_prompt[0, :-2])
substring_candidates = set()
restriction = min(len(substring_context) + 1, 16)
substring_context_splitted = substring_context.split(' ')
word_lengths = [len(word) for word in substring_context_splitted]
avg_length = (sum(word_lengths) + len(word_lengths))/len(word_lengths)
step_size = int(restriction // avg_length) + 1
# Iterate over the text to generate all possible substrings within the restriction
for i in range(len(substring_context_splitted)):
for j in range(i+1, min(len(substring_context_splitted) + 1, i + 1 + step_size)):
candidate = " ".join(substring_context_splitted[i:j])
substring_candidates.add(candidate)
# Sort the set of substring candidates for reproducibility
substring_candidates = sorted(substring_candidates)
# Prefix each candidate with the prompt to form complete sequences
substring_candidates = [last_part + candidate for candidate in substring_candidates]
tokenized_s_cand = tokenizer(substring_candidates,
return_tensors="pt",
padding=True,
add_special_tokens=False)['input_ids']
common_part_encoded = tokenizer.encode(common_part)
chosed_options = await fast_logprobs(common_part_encoded,
tokenized_s_cand,
1,
tokenizer)
res_text = []
for option in chosed_options:
substirng_choice = tokenizer.decode(option, skip_special_tokens=True)
res_text.append(substirng_choice[len(last_part):])
return JSONResponse(res_text)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--host", type=str, default=None)
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--ssl-keyfile", type=str, default=None)
parser.add_argument("--ssl-certfile", type=str, default=None)
parser = AsyncEngineArgs.add_cli_args(parser)
args = parser.parse_args()
engine_args: AsyncEngineArgs = AsyncEngineArgs.from_cli_args(args) # type: ignore
# Sets default for the model (`facebook/opt-125m`)
engine = AsyncLLMEngine.from_engine_args(engine_args)
uvicorn.run(
app,
host=args.host,
port=args.port,
log_level="debug",
timeout_keep_alive=TIMEOUT_KEEP_ALIVE,
ssl_keyfile=args.ssl_keyfile,
ssl_certfile=args.ssl_certfile,
)Request to the server code example:
import requests
import time
prompt = '''There is a request from the user to create a task in task management system. Please, extract the status of this task that will be placed in Task Management System.
Example 1: Update the status of the pentest task to completed.
Answer: “completed”
Example 2: “Create a code refactoring task for Vanya. And set the status to “in progress”
Answer: “in progress”.
Example 3: “Show me the executors of the task of charting”.
Answer: “null”.
Example 4: “Create a task to buy candy, status - in progress”.
Answer: “'''
substirng = "Create a task to buy candy, status - in progress"
substring_example = {'prompt': prompt, 'substring': substirng}
s = time.time()
x = requests.post(url, json = substring_example)
e = time.time()
print(x.text)
e - s
hmellor
robertgshaw2-neuralmagic
Delviet
Your current environment
🐛 Describe the bug
When attempting to enable the KV cache in the VLLM library by adding the parameter
--enable-prefix-cachingto the AsyncLLMEngine, I encountered the following exception message:This issue occurs despite following the documentation's instructions. The full output in the console: