Mirrors
/
hf_text-generation-inference
mirror of https://github.com/huggingface/text-generation-inference.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fixing linters. (#2650)

This commit is contained in:
Nicolas Patry 2024-10-15 12:43:49 +02:00 committed by GitHub
parent 58848cb471
commit cf04a43fb1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 4 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
server/text_generation_server/models/causal_lm.py
View File

@ -619,18 +619,11 @@ class CausalLM(Model):
model_id, model_id,
revision=revision, revision=revision,
torch_dtype=dtype, torch_dtype=dtype,
device_map=( device_map=("auto" if device_count > 1 else None),
"auto"
if device_count > 1
else None
),
load_in_8bit=quantize == "bitsandbytes", load_in_8bit=quantize == "bitsandbytes",
trust_remote_code=trust_remote_code, trust_remote_code=trust_remote_code,
) )
if ( if device_count == 1 and quantize != "bitsandbytes":
device_count == 1
and quantize != "bitsandbytes"
):
model = model.to(device) model = model.to(device)
if tokenizer.pad_token_id is None: if tokenizer.pad_token_id is None:

6
server/text_generation_server/models/seq2seq_lm.py
View File

@ -649,11 +649,7 @@ class Seq2SeqLM(Model):
model_id, model_id,
revision=revision, revision=revision,
torch_dtype=dtype, torch_dtype=dtype,
device_map=( device_map=("auto" if device_count > 1 else None),
"auto"
if device_count > 1
else None
),
load_in_8bit=quantize == "bitsandbytes", load_in_8bit=quantize == "bitsandbytes",
trust_remote_code=trust_remote_code, trust_remote_code=trust_remote_code,
) )