|
|
@@ -998,16 +998,20 @@ def generate_batched(
|
|
|
def generate_long_parallel_batched(
|
|
|
*,
|
|
|
model,
|
|
|
- device,
|
|
|
- decode_one_token,
|
|
|
- text,
|
|
|
- prompt_text=None,
|
|
|
- prompt_tokens=None,
|
|
|
- max_new_tokens=512,
|
|
|
- chunk_length=512,
|
|
|
- temperature=1.0,
|
|
|
- top_p=0.9,
|
|
|
- top_k=30,
|
|
|
+ device: Union[str, torch.device],
|
|
|
+ decode_one_token: Callable,
|
|
|
+ text: str,
|
|
|
+ num_samples: int = 1,
|
|
|
+ max_new_tokens: int = 0,
|
|
|
+ top_p: float = 0.9,
|
|
|
+ top_k: int = 30,
|
|
|
+ repetition_penalty: float = 1.1,
|
|
|
+ temperature: float = 1.0,
|
|
|
+ compile: bool = False,
|
|
|
+ iterative_prompt: bool = True,
|
|
|
+ chunk_length: int = 512,
|
|
|
+ prompt_text: Optional[Union[str, list[str]]] = None,
|
|
|
+ prompt_tokens: Optional[Union[torch.Tensor, list[torch.Tensor]]] = None,
|
|
|
):
|
|
|
"""
|
|
|
最小侵入版本:
|