From 21d5e51ba49269790ebd1db3fe1052683f42b019 Mon Sep 17 00:00:00 2001 From: Daniel Nichols Date: Tue, 29 Oct 2024 11:48:27 -0400 Subject: [PATCH] explicitly set the task in the HF pipeline so it can work in offline mode --- generate/generate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/generate/generate.py b/generate/generate.py index 6cb8e11..dfebd24 100644 --- a/generate/generate.py +++ b/generate/generate.py @@ -98,7 +98,7 @@ prompts_repeated = [p for p in prompts for _ in range(args.num_samples_per_prompt)] """ Initialize HuggingFace pipeline for generation """ -generator = pipeline(model=args.model, torch_dtype=inference_config.get_dtype(), device=0, token=args.hf_token) +generator = pipeline(task="text-generation", model=args.model, torch_dtype=inference_config.get_dtype(), device=0, token=args.hf_token) inference_config.init_padding(generator.tokenizer) """ Create a prompt data set to pass to generate method """