ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-01-16 01:40:33 +00:00 • committed by Daniel Han
commit 0b1bfaffcd
1 changed files with 3 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

12
unsloth/models/vision.py
View file

@ -296,9 +296,7 @@ def unsloth_base_fast_generate(
# Track metrics if enabled
collector = None
request_id = None
prompt_tokens_per_sequence = (
input_ids.shape[-1] if input_ids is not None else 0
)
prompt_tokens_per_sequence = input_ids.shape[-1] if input_ids is not None else 0
prompt_batch_size = input_ids.shape[0] if input_ids.dim() > 1 else 1
num_prompt_tokens = prompt_tokens_per_sequence * prompt_batch_size
max_tokens = kwargs.get("max_new_tokens") or kwargs.get("max_length")
@ -359,9 +357,7 @@ def unsloth_base_fast_generate(
# Handle ModelOutput when return_dict_in_generate=True
sequences = output["sequences"]
if isinstance(sequences, torch.Tensor):
output_batch_size = (
sequences.shape[0] if sequences.dim() > 1 else 1
)
output_batch_size = sequences.shape[0] if sequences.dim() > 1 else 1
if output_batch_size != prompt_batch_size:
effective_prompt_tokens = (
prompt_tokens_per_sequence * output_batch_size
@ -377,9 +373,7 @@ def unsloth_base_fast_generate(
# Handle ModelOutput object directly
sequences = output.sequences
if isinstance(sequences, torch.Tensor):
output_batch_size = (
sequences.shape[0] if sequences.dim() > 1 else 1
)
output_batch_size = sequences.shape[0] if sequences.dim() > 1 else 1
if output_batch_size != prompt_batch_size:
effective_prompt_tokens = (
prompt_tokens_per_sequence * output_batch_size