mirror of
https://github.com/hiyouga/LLaMA-Factory.git
synced 2025-08-06 13:42:51 +08:00
91 lines
3.1 KiB
Python
91 lines
3.1 KiB
Python
import asyncio
|
|
from threading import Thread
|
|
from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, Generator, List, Optional, Sequence
|
|
|
|
from ..hparams import get_infer_args
|
|
from .hf_engine import HuggingfaceEngine
|
|
from .vllm_engine import VllmEngine
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from .base_engine import BaseEngine, Response
|
|
|
|
|
|
def _start_background_loop(loop: asyncio.AbstractEventLoop) -> None:
|
|
asyncio.set_event_loop(loop)
|
|
loop.run_forever()
|
|
|
|
|
|
class ChatModel:
|
|
def __init__(self, args: Optional[Dict[str, Any]] = None) -> None:
|
|
model_args, data_args, finetuning_args, generating_args = get_infer_args(args)
|
|
if model_args.infer_backend == "huggingface":
|
|
self.engine: "BaseEngine" = HuggingfaceEngine(model_args, data_args, finetuning_args, generating_args)
|
|
elif model_args.infer_backend == "vllm":
|
|
self.engine: "BaseEngine" = VllmEngine(model_args, data_args, finetuning_args, generating_args)
|
|
else:
|
|
raise NotImplementedError("Unknown backend: {}".format(model_args.infer_backend))
|
|
|
|
self._loop = asyncio.new_event_loop()
|
|
self._thread = Thread(target=_start_background_loop, args=(self._loop,), daemon=True)
|
|
self._thread.start()
|
|
|
|
def chat(
|
|
self,
|
|
messages: Sequence[Dict[str, str]],
|
|
system: Optional[str] = None,
|
|
tools: Optional[str] = None,
|
|
**input_kwargs,
|
|
) -> List["Response"]:
|
|
task = asyncio.run_coroutine_threadsafe(self.achat(messages, system, tools, **input_kwargs), self._loop)
|
|
return task.result()
|
|
|
|
async def achat(
|
|
self,
|
|
messages: Sequence[Dict[str, str]],
|
|
system: Optional[str] = None,
|
|
tools: Optional[str] = None,
|
|
**input_kwargs,
|
|
) -> List["Response"]:
|
|
return await self.engine.chat(messages, system, tools, **input_kwargs)
|
|
|
|
def stream_chat(
|
|
self,
|
|
messages: Sequence[Dict[str, str]],
|
|
system: Optional[str] = None,
|
|
tools: Optional[str] = None,
|
|
**input_kwargs,
|
|
) -> Generator[str, None, None]:
|
|
generator = self.astream_chat(messages, system, tools, **input_kwargs)
|
|
while True:
|
|
try:
|
|
task = asyncio.run_coroutine_threadsafe(generator.__anext__(), self._loop)
|
|
yield task.result()
|
|
except StopAsyncIteration:
|
|
break
|
|
|
|
async def astream_chat(
|
|
self,
|
|
messages: Sequence[Dict[str, str]],
|
|
system: Optional[str] = None,
|
|
tools: Optional[str] = None,
|
|
**input_kwargs,
|
|
) -> AsyncGenerator[str, None]:
|
|
async for new_token in self.engine.stream_chat(messages, system, tools, **input_kwargs):
|
|
yield new_token
|
|
|
|
def get_scores(
|
|
self,
|
|
batch_input: List[str],
|
|
**input_kwargs,
|
|
) -> List[float]:
|
|
task = asyncio.run_coroutine_threadsafe(self.aget_scores(batch_input, **input_kwargs), self._loop)
|
|
return task.result()
|
|
|
|
async def aget_scores(
|
|
self,
|
|
batch_input: List[str],
|
|
**input_kwargs,
|
|
) -> List[float]:
|
|
return await self.engine.get_scores(batch_input, **input_kwargs)
|