From febc4aaf82cad76685b17745823a553d1c7ff482 Mon Sep 17 00:00:00 2001
From: Vela-zz <59645193+Vela-zz@users.noreply.github.com>
Date: Sat, 2 Dec 2023 02:08:04 +0900
Subject: [PATCH] Update detoxify.py

align action with transformers
---
 detoxify/detoxify.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/detoxify/detoxify.py b/detoxify/detoxify.py
index 4a25dca..5837348 100644
--- a/detoxify/detoxify.py
+++ b/detoxify/detoxify.py
@@ -1,6 +1,5 @@
 import torch
 import transformers
-from transformers import AutoConfig
 
 
 DOWNLOAD_URL = "https://github.com/unitaryai/detoxify/releases/download/"
@@ -19,8 +18,7 @@ def get_model_and_tokenizer(
     model_type, model_name, tokenizer_name, num_classes, state_dict, huggingface_config_path=None
 ):
     model_class = getattr(transformers, model_name)
-    config = AutoConfig.from_pretrained(model_type)
-    config.num_labels = num_classes
+    config = model_class.config_class.from_pretrained(model_type, num_labels=num_classes)
     model = model_class.from_pretrained(
         pretrained_model_name_or_path=None,
         config=huggingface_config_path or config,