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

Update vision.py

Browse source
This commit is contained in:
Daniel Han 2025-05-16 23:20:18 -07:00
commit caae7668e0
1 changed files with 7 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

9
unsloth/models/vision.py
View file

@ -89,15 +89,20 @@ def unsloth_base_fast_generate(
*args,
**kwargs,
):
print(args, kwargs)
if len(args) != 0:
input_ids = args[0]
elif "input_ids" in kwargs:
input_ids = kwargs["input_ids"]
elif "input" in kwargs:
input_ids = kwargs["input_ids"]
elif "inputs" in kwargs:
input_ids = kwargs["inputs"]
else:
raise TypeError("Unsloth: You need to pass in input_ids to .generate!")
key = next(iter(kwargs.keys()))
if type(kwargs["key"]) is not torch.Tensor:
raise TypeError("Unsloth: You need to pass in input_ids to .generate!")
input_ids = kwargs[key]
pass
assert(type(input_ids) is torch.Tensor)
bsz = input_ids.shape[0]