ayo/ipex-llm
1
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

LLM: support tensor input of native int4 generate (#8620)

Browse source
This commit is contained in:
binbin Deng 2023-07-27 17:59:49 +08:00 committed by GitHub
parent 5b484ab48d
commit fb32fefcbe
1 changed files with 9 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
python/llm/src/bigdl/llm/ggml/model/generation/utils.py
View file

@ -22,6 +22,7 @@
from typing import Optional, Union, Sequence, List from typing import Optional, Union, Sequence, List
from bigdl.llm.utils.common import invalidInputError from bigdl.llm.utils.common import invalidInputError
import torch
class GenerationMixin: class GenerationMixin:
@ -100,8 +101,9 @@ class GenerationMixin:
def generate( def generate(
self, self,
inputs: Union[Optional[Sequence[int]], inputs: Optional[Union[Sequence[int],
Sequence[Sequence[int]]]=None, Sequence[Sequence[int]],
torch.Tensor]]=None,
max_new_tokens: int = 128, max_new_tokens: int = 128,
top_k: int = 40, top_k: int = 40,
top_p: float = 0.95, top_p: float = 0.95,
@ -116,9 +118,9 @@ class GenerationMixin:
mirostat_eta: float = 0.1, mirostat_eta: float = 0.1,
stop: Optional[Union[str, List[str]]]=[], # TODO: rebase to support stopping_criteria stop: Optional[Union[str, List[str]]]=[], # TODO: rebase to support stopping_criteria
**kwargs, **kwargs,
) -> Union[Optional[Sequence[int]], ) -> Optional[Union[Sequence[int],
Sequence[Sequence[int]], Sequence[Sequence[int]],
None]: None]]:
# TODO: modify docs # TODO: modify docs
"""Create a generator of tokens from a prompt. """Create a generator of tokens from a prompt.
@ -140,6 +142,8 @@ class GenerationMixin:
Yields: Yields:
The generated tokens. The generated tokens.
""" """
if isinstance(inputs, torch.Tensor):
inputs = inputs.tolist()
if inputs and len(inputs) > 0: if inputs and len(inputs) > 0:
if not isinstance(inputs[0], Sequence): if not isinstance(inputs[0], Sequence):
inputs = [inputs] inputs = [inputs]