/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
colossalai/inference/server/completion_service.py
34 строки
1 KB
Jianghai
[Inference] Fix API server, test and example (#5712)
15 май 2024, 10:47
Не верифицирован
15 май 2024, 10:47
f47f2fb
Код
Авторство
О чём код?
import asyncio from colossalai.inference.core.async_engine import AsyncInferenceEngine from .utils import id_generator class CompletionServing: def __init__(self, engine: AsyncInferenceEngine, served_model: str): self.engine = engine self.served_model = served_model try: asyncio.get_running_loop() except RuntimeError: pass async def create_completion(self, request, generation_config): request_dict = await request.json() request_id = id_generator() prompt = request_dict.pop("prompt") # it is not a intuitive way self.engine.engine.generation_config = generation_config result_generator = self.engine.generate(request_id, prompt=prompt, generation_config=generation_config) if await request.is_disconnected(): # Abort the request if the client disconnects. await self.engine.abort(request_id) raise RuntimeError("Client disconnected") final_res = await result_generator return final_res