Skip to content

Commit

Permalink
review(AutoGenerator): log if using Jetstream/PT or torch xla
Browse files Browse the repository at this point in the history
  • Loading branch information
tengomucho committed Sep 6, 2024
1 parent eb677bd commit 1f5e9c4
Showing 1 changed file with 4 additions and 1 deletion.
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from .generator_base import Generator
from .jetstream_pt_support import model_can_use_jetstream_pt
from loguru import logger


class AutoGenerator:
Expand All @@ -23,12 +24,14 @@ def from_pretrained(
Returns:
A TpuGenerator.
"""
if check(model_path):
if model_can_use_jetstream_pt(model_path):
logger.debug("Using Jetstream PyTorch generator.")
from .jetstream_pt_support.generator import TpuGeneratorJetStream
return TpuGeneratorJetStream.from_pretrained(
model_path, revision=revision, max_batch_size=max_batch_size, max_sequence_length=max_sequence_length
)
else:
logger.debug("Using PyTorch/XLA generator.")
from .generator import TpuGenerator
return TpuGenerator.from_pretrained(
model_path, revision=revision, max_batch_size=max_batch_size, max_sequence_length=max_sequence_length
Expand Down

0 comments on commit 1f5e9c4

Please sign in to comment.