Repository: huggingface/text-generation-inference

Issue:
No trust_remote_code passed to AutoTokenizer
### System Info

There is no trust_remote_code option passed in to the `AutoTokenizer` of `get_wikitext2` function (https://github.com/huggingface/text-generation-inference/blob/main/server/text_generation_server/utils/gptq/quantize.py#L363). This causes a `trust_remove_code=False` issue when I try to load a base model and quantize it with GPTQ with the command below

```
docker run \
    -it \
    --gpus all \
    --shm-size 1g \
    --pull always \
    -v $volume:/data \
    --entrypoint text-generation-server \
    ghcr.io/huggingface/text-generation-inference:latest \
    quantize \
    --trust-remote-code \
    /data/$model \
    /data/$model-quantized
```

The current form of `get_wikitext2` looks like below:
```
def get_wikitext2(nsamples, seed, seqlen, model_id):
    from datasets import load_dataset

    traindata = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
    testdata = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")

    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
    trainenc = tokenizer("\n\n".join(traindata["text"]), return_tensors="pt")
    testenc = tokenizer("\n\n".join(testdata["text"]), return_tensors="pt")

    import random

    random.seed(seed)
    trainloader = []
    for _ in range(nsamples):
        i = random.randint(0, trainenc.input_ids.shape[1] - seqlen - 1)
        j = i + seqlen
        inp = trainenc.input_ids[:, i:j]
        tar = inp.clone()
        tar[:, :-1] = -100
        trainloader.append((inp, tar))
    return trainloader, testenc
```

### Information

- [ ] Docker
- [X] The CLI directly

### Tasks

- [X] An officially supported command
- [ ] My own modifications

### Reproduction

1. docker run \
    -it \
    --gpus all \
    --shm-size 1g \
    --pull always \
    -v $volume:/data \
    --entrypoint text-generation-server \
    ghcr.io/huggingface/text-generation-inference:latest \
    quantize \
    --trust-remote-code \
    /data/$model \
    /data/$model-quantized

2. See error describe above related to `trust_remote_code=False`

### Expected behavior

I would expect the specified `trust_remote_code` option is passed into this function.


Discussion:

Resolving patch:
diff --git a/server/text_generation_server/utils/gptq/quantize.py b/server/text_generation_server/utils/gptq/quantize.py
--- a/server/text_generation_server/utils/gptq/quantize.py
+++ b/server/text_generation_server/utils/gptq/quantize.py
@@ -360,15 +360,21 @@ def free(self):
         torch.cuda.empty_cache()
 
 
-def get_wikitext2(nsamples, seed, seqlen, model_id):
+def get_wikitext2(nsamples, seed, seqlen, model_id, trust_remote_code):
     from datasets import load_dataset
 
     traindata = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
     testdata = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
 
-    from transformers import AutoTokenizer
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_id, use_fast=False, trust_remote_code=trust_remote_code
+        )
+    except:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_id, use_fast=True, trust_remote_code=trust_remote_code
+        )
 
-    tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
     trainenc = tokenizer("\n\n".join(traindata["text"]), return_tensors="pt")
     testenc = tokenizer("\n\n".join(testdata["text"]), return_tensors="pt")
 
@@ -386,18 +392,21 @@ def get_wikitext2(nsamples, seed, seqlen, model_id):
     return trainloader, testenc
 
 
-def get_ptb(nsamples, seed, seqlen, model_id):
+def get_ptb(nsamples, seed, seqlen, model_id, trust_remote_code):
     from datasets import load_dataset
 
     traindata = load_dataset("ptb_text_only", "penn_treebank", split="train")
     valdata = load_dataset("ptb_text_only", "penn_treebank", split="validation")
 
-    from transformers import AutoTokenizer
-
     try:
-        tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_id, use_fast=False, trust_remote_code=trust_remote_code
+        )
     except:
-        tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_id, use_fast=True, trust_remote_code=trust_remote_code
+        )
+
     trainenc = tokenizer("\n\n".join(traindata["sentence"]), return_tensors="pt")
     testenc = tokenizer("\n\n".join(valdata["sentence"]), return_tensors="pt")
 
@@ -415,7 +424,7 @@ def get_ptb(nsamples, seed, seqlen, model_id):
     return trainloader, testenc
 
 
-def get_c4(nsamples, seed, seqlen, model_id):
+def get_c4(nsamples, seed, seqlen, model_id, trust_remote_code):
     from datasets import load_dataset
 
     traindata = load_dataset(
@@ -433,12 +442,14 @@ def get_c4(nsamples, seed, seqlen, model_id):
         use_auth_token=False,
     )
 
-    from transformers import AutoTokenizer
-
     try:
-        tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_id, use_fast=False, trust_remote_code=trust_remote_code
+        )
     except:
-        tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_id, use_fast=True, trust_remote_code=trust_remote_code
+        )
 
     import random
 
@@ -481,18 +492,21 @@ def __init__(self, input_ids):
     return trainloader, valenc
 
 
-def get_ptb_new(nsamples, seed, seqlen, model_id):
+def get_ptb_new(nsamples, seed, seqlen, model_id, trust_remote_code):
     from datasets import load_dataset
 
     traindata = load_dataset("ptb_text_only", "penn_treebank", split="train")
     testdata = load_dataset("ptb_text_only", "penn_treebank", split="test")
 
-    from transformers import AutoTokenizer
-
     try:
-        tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_id, use_fast=False, trust_remote_code=trust_remote_code
+        )
     except:
-        tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_id, use_fast=True, trust_remote_code=trust_remote_code
+        )
+
     trainenc = tokenizer(" ".join(traindata["sentence"]), return_tensors="pt")
     testenc = tokenizer(" ".join(testdata["sentence"]), return_tensors="pt")
 
@@ -510,7 +524,7 @@ def get_ptb_new(nsamples, seed, seqlen, model_id):
     return trainloader, testenc
 
 
-def get_c4_new(nsamples, seed, seqlen, model_id):
+def get_c4_new(nsamples, seed, seqlen, model_id, trust_remote_code):
     from datasets import load_dataset
 
     traindata = load_dataset(
@@ -526,12 +540,14 @@ def get_c4_new(nsamples, seed, seqlen, model_id):
         split="validation",
     )
 
-    from transformers import AutoTokenizer
-
     try:
-        tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_id, use_fast=False, trust_remote_code=trust_remote_code
+        )
     except:
-        tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_id, use_fast=True, trust_remote_code=trust_remote_code
+        )
 
     import random
 
@@ -562,17 +578,17 @@ def __init__(self, input_ids):
     return trainloader, valenc
 
 
-def get_loaders(name, nsamples=128, seed=0, seqlen=2048, model_id=""):
+def get_loaders(name, nsamples=128, seed=0, seqlen=2048, model_id="", trust_remote_code=False):
     if "wikitext2" in name:
-        return get_wikitext2(nsamples, seed, seqlen, model_id)
+        return get_wikitext2(nsamples, seed, seqlen, model_id, trust_remote_code)
     if "ptb" in name:
         if "new" in name:
-            return get_ptb_new(nsamples, seed, seqlen, model_id)
-        return get_ptb(nsamples, seed, seqlen, model_id)
+            return get_ptb_new(nsamples, seed, seqlen, model_id, trust_remote_code)
+        return get_ptb(nsamples, seed, seqlen, model_id, trust_remote_code)
     if "c4" in name:
         if "new" in name:
-            return get_c4_new(nsamples, seed, seqlen, model_id)
-        return get_c4(nsamples, seed, seqlen, model_id)
+            return get_c4_new(nsamples, seed, seqlen, model_id, trust_remote_code)
+        return get_c4(nsamples, seed, seqlen, model_id, trust_remote_code)
 
 
 def find_layers(module, layers=(nn.Conv2d, nn.Linear), name=""):
@@ -906,7 +922,12 @@ def _unload():
     seed = None
 
     dataloader, testloader = get_loaders(
-        dataset, nsamples=nsamples, seed=seed, model_id=model_id, seqlen=model.seqlen
+        dataset,
+        nsamples=nsamples,
+        seed=seed,
+        model_id=model_id,
+        seqlen=model.seqlen,
+        trust_remote_code=trust_remote_code
     )
 
     tick = time.time()


Test patch: