From 76fbf94310272b53932a804234d3f66fe853e85f Mon Sep 17 00:00:00 2001 From: Alvaro Moran Date: Fri, 6 Sep 2024 09:35:35 +0000 Subject: [PATCH] review(AutoGenerator): log if using Jetstream/PT or torch xla --- .../server/text_generation_server/auto_generator.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/text-generation-inference/server/text_generation_server/auto_generator.py b/text-generation-inference/server/text_generation_server/auto_generator.py index 9a549451..23e5631b 100644 --- a/text-generation-inference/server/text_generation_server/auto_generator.py +++ b/text-generation-inference/server/text_generation_server/auto_generator.py @@ -1,3 +1,5 @@ +from loguru import logger + from .generator_base import Generator from .jetstream_pt_support import model_can_use_jetstream_pt @@ -23,12 +25,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