🐛 Fix trust_remote_code=True bug in tests

Signed-off-by: Marcos Martínez Galindo <marcosmartinezgalindo@marcoss-mbp.mul.ie.ibm.com>
IBM · Jul 23, 2024 · 6a48d07 · 6a48d07
1 parent 5bc52a5
commit 6a48d07
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/zshot/evaluation/dataset/fewrel/fewrel.py b/zshot/evaluation/dataset/fewrel/fewrel.py
@@ -25,7 +25,7 @@ def get_entity_data(e, tokenized_sentence):
 
 def load_few_rel_zs(split: Optional[Union[str, Split]] = "val_wiki") -> Union[Dict[DatasetWithRelations,
                                                                                    Dataset], Dataset]:
-    dataset = load_dataset("few_rel", split=split)
+    dataset = load_dataset("few_rel", split=split, trust_remote_code=True)
     relations_descriptions = dataset["names"]
     tokenized_sentences = dataset["tokens"]
     sentences = [" ".join(tokens) for tokens in tokenized_sentences]

diff --git a/zshot/evaluation/dataset/med_mentions/med_mentions.py b/zshot/evaluation/dataset/med_mentions/med_mentions.py
@@ -13,7 +13,7 @@
 
 
 def load_medmentions_zs(split: Optional[Union[str, Split]] = None, **kwargs) -> Union[DatasetDict, Dataset]:
-    dataset = load_dataset(REPO_ID, split=split, **kwargs)
+    dataset = load_dataset(REPO_ID, split=split, trust_remote_code=True, **kwargs)
     entities_file = hf_hub_download(repo_id=REPO_ID,
                                     repo_type='dataset',
                                     filename=ENTITIES_FN)

diff --git a/zshot/evaluation/dataset/ontonotes/onto_notes.py b/zshot/evaluation/dataset/ontonotes/onto_notes.py
@@ -58,7 +58,7 @@ def remove_out_of_split(sentence, split):
 def load_ontonotes_zs(split: Optional[Union[str, Split]] = None, **kwargs) -> Union[Dict[DatasetWithEntities,
                                                                                          Dataset], Dataset]:
     dataset_zs = load_dataset("conll2012_ontonotesv5", "english_v12",
-                              split=split, verification_mode='no_checks', **kwargs)
+                              split=split, verification_mode='no_checks', trust_remote_code=True, **kwargs)
     if split:
         ontonotes_zs = preprocess_spit(dataset_zs, get_simple_split(split))
     else: