fix trust_remote_code

Author: tastelikefeet

Dockerfile

@@ -21,6 +21,3 @@ RUN sh INSTALL_MEGATRON.sh
 RUN pip install --no-cache-dir tinker==0.14.0 "ray[serve]" transformers peft accelerate -U
 
 RUN pip install -e . --no-build-isolation
-
-ENV TWINKLE_WORKDIR=/data
-CMD ["bash", "-c", "mkdir -p $TWINKLE_WORKDIR && cd $TWINKLE_WORKDIR && bash /twinkle/cookbook/client/server/megatron/run.sh 2>&1 | tee $TWINKLE_WORKDIR/run.log"]

cookbook/client/server/megatron/run.sh

@@ -3,4 +3,4 @@ export RAY_ROTATION_BACKUP_COUNT=1
 CUDA_VISIBLE_DEVICES=0,1,2,3 ray start --head --port=6379 --num-gpus=4 --disable-usage-stats --include-dashboard=false
 CUDA_VISIBLE_DEVICES=4,5,6,7 ray start --address=127.0.0.1:6379 --num-gpus=4
 CUDA_VISIBLE_DEVICES="" ray start --address=127.0.0.1:6379 --num-gpus=0
-python server.py
+python "$(dirname "$0")/server.py"

cookbook/client/server/megatron/server.py

@@ -9,7 +9,7 @@
 import os
 
 # Enable Ray debug mode for verbose logging during development
-os.environ['TWINKLE_TRUST_REMOTE_CODE'] = '1'
+os.environ['TWINKLE_TRUST_REMOTE_CODE'] = '0'
 
 from twinkle.server import launch_server
 

src/twinkle/dataset/base.py

@@ -51,6 +51,9 @@ class Dataset(TorchDataset):
     """
 
     def __init__(self, dataset_meta: DatasetMeta, **kwargs):
+        trust_remote_code = bool(os.environ.get('TWINKLE_TRUST_REMOTE_CODE', '1'))
+        if not trust_remote_code:
+            kwargs['trust_remote_code'] = False
         dataset = self._load_dataset(dataset_meta, **kwargs)
         self.datasets = {dataset_meta.get_id(): dataset}
         self.dataset = dataset
@@ -247,6 +250,9 @@ def add_dataset(self, dataset_meta: DatasetMeta, **kwargs):
         Args:
             dataset_meta: The dataset_meta information of the loaded dataset.
         """
+        trust_remote_code = bool(os.environ.get('TWINKLE_TRUST_REMOTE_CODE', '1'))
+        if not trust_remote_code:
+            kwargs['trust_remote_code'] = False
         dataset = self._load_dataset(dataset_meta, **kwargs)
         self.datasets[dataset_meta.get_id()] = dataset
 

src/twinkle/hub/hub.py

@@ -401,7 +401,7 @@ def load_dataset(cls,
         cls.try_login(token)
         if revision is None or revision == 'main':
             revision = 'master'
-        load_kwargs = {'trust_remote_code': True}
+        load_kwargs = {'trust_remote_code': kwargs.get('trust_remote_code', True)}
         return MsDataset.load(
             dataset_id,
             subset_name=subset_name,
@@ -595,13 +595,15 @@ def load_dataset(cls,
         from datasets import load_dataset
         if revision is None or revision == 'master':
             revision = 'main'
+        trust_remote_code = kwargs.get('trust_remote_code', True)
         return load_dataset(
             dataset_id,
             name=subset_name,
             split=split,
             streaming=streaming,
             revision=revision,
             download_mode=download_mode,
+            trust_remote_code=trust_remote_code,
             num_proc=num_proc)
 
     @classmethod