haowei2000 · deepsource-autofix · Jan 1, 2025
diff --git a/src/everyai/classifier/classify.py b/src/everyai/classifier/classify.py
@@ -43,31 +43,26 @@ def label_encode(labels):
     return LabelEncoder().fit_transform(labels)
 
 
-def split_data(x, y, train_size=0.8, valid_size=0.1, test_size=0.1):# -> tuple:
+# -> tuple:
+def split_data(x, y, train_size=0.8, valid_size=0.1, test_size=0.1):
     # 获取原始数据的索引
-    original_indices = (
-        x.index if isinstance(x, pd.DataFrame) else np.arange(x.shape[0])
-    )
+    original_indices = x.index if isinstance(x, pd.DataFrame) else np.arange(x.shape[0])
 
     # 第一次划分: 训练集和测试集
-    x_train, x_test, y_train, y_test, train_indices, test_indices = (
-        train_test_split(
-            x,
-            y,
-            original_indices,
-            test_size=test_size,
-            random_state=42,
-        )
+    x_train, x_test, y_train, y_test, train_indices, test_indices = train_test_split(
+        x,
+        y,
+        original_indices,
+        test_size=test_size,
+        random_state=42,
     )
     # 第二次划分: 训练集和验证集
-    x_train, x_valid, y_train, y_valid, train_indices, valid_indices = (
-        train_test_split(
-            x_train,
-            y_train,
-            train_indices,
-            test_size=valid_size / (train_size + valid_size),
-            random_state=42,
-        )
+    x_train, x_valid, y_train, y_valid, train_indices, valid_indices = train_test_split(
+        x_train,
+        y_train,
+        train_indices,
+        test_size=valid_size / (train_size + valid_size),
+        random_state=42,
     )
     return (
         x_train,
@@ -211,8 +206,7 @@ def __init__(
         )
         self.model_config = None
         self.model_path = (
-            MODEL_PATH
-            / f"{self.model_name}_{self.tokenizer_name}_{self.data_name}.pkl"
+            MODEL_PATH / f"{self.model_name}_{self.tokenizer_name}_{self.data_name}.pkl"
         )
         self.pipeline = pipeline if pipeline is not None else None
 
@@ -222,9 +216,7 @@ def load_data(self, texts, labels, data_name):
             raise ValueError("Length of texts and labels should be same")
         self.texts = texts
         self.labels = labels
-        logging.info(
-            "Loading data: %s to classfier %s", data_name, self.model_name
-        )
+        logging.info("Loading data: %s to classfier %s", data_name, self.model_name)
         self.data_name = data_name
         self.classfier_name = (
             f"{self.model_name}_{self.tokenizer_name}_{self.data_name}"

diff --git a/src/everyai/classifier/huggingface_classifier.py b/src/everyai/classifier/huggingface_classifier.py
@@ -1,9 +1,10 @@
 from typing import List
 
-from sklearn.calibration import LabelEncoder
 import torch
+from sklearn.calibration import LabelEncoder
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
 from everyai.classifier.classify import TextClassifer
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
 
 class HuggingfaceClassifer(TextClassifer):
@@ -31,23 +32,21 @@ def __init__(
             tokenizer,
         )
         self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
-        self.model = AutoModelForSequenceClassification.from_pretrained(
-            model_name
-        )
-        self.tokenzier_config  = tokenzier_config
-
+        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
+        self.tokenzier_config = tokenzier_config
+
     def _label_encode(self, labels: List[str]):
         label_encoder = LabelEncoder()
         label_ids = label_encoder.fit_transform(labels)
         return torch.tensor(label_ids, dtype=torch.long)
 
-    def _tokenize(self, texts,labels):
+    def _tokenize(self, texts, labels):
         def tokenzier_fn(examples):
             return self.tokenizer(texts, **self.tokenzier_config)
-    
+
     def train(self):
         self.data.y = self._label_encode(self.labels)
         self.model.train()
         self.model.to(self.device)
         self.model.train()
-        return self.model
+        return self.model
diff --git a/src/everyai/classifier/sklearn_classifier.py b/src/everyai/classifier/sklearn_classifier.py
@@ -39,9 +39,7 @@ def _init_sklearn_pipeline(pipeline_config: list[dict]):
             step_params = pipeline_config[step_name]
             steps.append((step_name, step_dict[step_name](**step_params)))
         else:
-            logging.warning(
-                "Step %s not recognized and will be skipped", step_name
-            )
+            logging.warning("Step %s not recognized and will be skipped", step_name)
     return make_pipeline(*[step[1] for step in steps])
 
 
@@ -74,9 +72,7 @@ def __init__(self, **classfiy_config):
             raise ValueError("Model not supported")
 
         if classfiy_config["tokenizer_name"] in tokenizer_dict:
-            self.tokenizer = tokenizer_dict[
-                classfiy_config["tokenizer_name"]
-            ]()
+            self.tokenizer = tokenizer_dict[classfiy_config["tokenizer_name"]]()
         else:
             logging.error("Tokenizer not supported")
             raise ValueError("Tokenizer not supported")

diff --git a/src/everyai/config/config.py b/src/everyai/config/config.py
@@ -2,9 +2,13 @@
 
 import yaml
 
-from everyai.everyai_path import (BERT_TOPIC_CONFIG_PATH, CLASSIFY_CONFIG_PATH,
-                                  DATA_LOAD_CONFIG_PATH, GENERATE_CONFIG_PATH,
-                                  MONGO_CONFIG_PATH)
+from everyai.everyai_path import (
+    BERT_TOPIC_CONFIG_PATH,
+    CLASSIFY_CONFIG_PATH,
+    DATA_LOAD_CONFIG_PATH,
+    GENERATE_CONFIG_PATH,
+    MONGO_CONFIG_PATH,
+)
 
 
 def get_config(file_path: Path):

diff --git a/src/everyai/main.py b/src/everyai/main.py
@@ -8,9 +8,15 @@
 from everyai.data_loader.dataprocess import split_remove_stopwords_punctuation
 from everyai.data_loader.everyai_dataset import EveryaiDataset
 from everyai.data_loader.mongo_connection import get_mongo_connection
-from everyai.everyai_path import (BERT_TOPIC_CONFIG_PATH, CLASSIFY_CONFIG_PATH,
-                                  DATA_LOAD_CONFIG_PATH, DATA_PATH, FIG_PATH,
-                                  GENERATE_CONFIG_PATH, MONGO_CONFIG_PATH)
+from everyai.everyai_path import (
+    BERT_TOPIC_CONFIG_PATH,
+    CLASSIFY_CONFIG_PATH,
+    DATA_LOAD_CONFIG_PATH,
+    DATA_PATH,
+    FIG_PATH,
+    GENERATE_CONFIG_PATH,
+    MONGO_CONFIG_PATH,
+)
 from everyai.explanation.explain import LimeExplanation, ShapExplanation
 from everyai.generator.generate import Generator
 from everyai.topic.my_bertopic import create_topic
@@ -33,16 +39,12 @@ def generate():
                 file_path=data_config["file_path"],
                 data_type=data_config["data_type"],
             )
-            qa_datas = data_loader.load_data2list(
-                max_count=data_config["max_count"]
-            )
+            qa_datas = data_loader.load_data2list(max_count=data_config["max_count"])
             everyai_dataset = EveryaiDataset(
                 dataname=data_config["data_name"],
                 ai_list=[generate_config["model_name"]],
             )
-            for data in tqdm(
-                qa_datas, desc="Generating data", total=len(qa_datas)
-            ):
+            for data in tqdm(qa_datas, desc="Generating data", total=len(qa_datas)):
                 ai_response: str = generator.generate(data["question"])
                 everyai_dataset.insert_ai_response(
                     question=data["question"],
@@ -100,9 +102,7 @@ def classify():
         )
         everyai_dataset.load(formatter="mongodb")
         logging.info("Loaded data: %s", everyai_dataset.data_name)
-        texts, labels = everyai_dataset.get_records_with_1ai(
-            ["THUDM/glm-4-9b-chat-hf"]
-        )
+        texts, labels = everyai_dataset.get_records_with_1ai(["THUDM/glm-4-9b-chat-hf"])
         for classify_config in get_config(file_path=CLASSIFY_CONFIG_PATH)[
             "classifier_list"
         ]: