diff --git a/template/steps/dataset_loader/data_loader.py b/template/steps/dataset_loader/data_loader.py index 9d5fb9f..35ff1e6 100644 --- a/template/steps/dataset_loader/data_loader.py +++ b/template/steps/dataset_loader/data_loader.py @@ -32,14 +32,23 @@ def data_loader( # Load dataset based on the dataset value {%- if dataset == 'financial_news' %} - dataset = load_dataset("zeroshot/twitter-financial-news-sentiment") + dataset = load_dataset( + "zeroshot/twitter-financial-news-sentiment", + trust_remote_code=True, + ) {%- endif %} {%- if dataset == 'imdb_reviews' %} - dataset = load_dataset("imdb")["train"] + dataset = load_dataset( + "imdb", + trust_remote_code=True, + )["train"] dataset = dataset.train_test_split(test_size=0.25, shuffle=True) {%- endif %} {%- if dataset == 'airline_reviews' %} - dataset = load_dataset("Shayanvsf/US_Airline_Sentiment") + dataset = load_dataset( + "Shayanvsf/US_Airline_Sentiment", + trust_remote_code=True, + ) dataset = dataset.rename_column("airline_sentiment", "label") dataset = dataset.remove_columns(["airline_sentiment_confidence","negativereason_confidence"]) {%- endif %} diff --git a/template/utils/misc.py b/template/utils/misc.py index f8d326f..a133a73 100644 --- a/template/utils/misc.py +++ b/template/utils/misc.py @@ -18,13 +18,13 @@ def compute_metrics(eval_pred: Tuple[np.ndarray, np.ndarray]) -> Dict[str, float logits, labels = eval_pred predictions = np.argmax(logits, axis=-1) # calculate the mertic using the predicted and true value - accuracy = load_metric("accuracy").compute( + accuracy = load_metric("accuracy", trust_remote_code=True).compute( predictions=predictions, references=labels ) - f1 = load_metric("f1").compute( + f1 = load_metric("f1", trust_remote_code=True).compute( predictions=predictions, references=labels, average="weighted" ) - precision = load_metric("precision").compute( + precision = load_metric("precision", trust_remote_code=True).compute( predictions=predictions, references=labels, average="weighted" ) return {