diff --git a/src/ai_detection/classifier/classify.py b/src/ai_detection/classifier/classify.py index 4129053..ad47c28 100644 --- a/src/ai_detection/classifier/classify.py +++ b/src/ai_detection/classifier/classify.py @@ -6,7 +6,8 @@ from sklearn.calibration import LabelEncoder from sklearn.model_selection import train_test_split -from ai_detection.data_loader.data_process import split_remove_stopwords_punctuation +from ai_detection.data_loader.data_process import \ + split_remove_stopwords_punctuation from ai_detection.utils.everyai_path import MODEL_PATH from ai_detection.utils.load_args import set_attrs_2class @@ -50,23 +51,19 @@ def split_data( """ assert len(x) == len(y), "Length of x and y should be same" original_indices = pd.DataFrame(x).index - x_train, x_test, y_train, y_test, indices_train, indices_test = ( - train_test_split( - x, - y, - original_indices, - train_size=train_size, - random_state=42, - ) + x_train, x_test, y_train, y_test, indices_train, indices_test = train_test_split( + x, + y, + original_indices, + train_size=train_size, + random_state=42, ) - x_test, x_valid, y_test, y_valid, indices_train, indices_valid = ( - train_test_split( - x_test, - y_test, - indices_test, - test_size=test_size / (test_size + valid_size), - random_state=42, - ) + x_test, x_valid, y_test, y_valid, indices_train, indices_valid = train_test_split( + x_test, + y_test, + indices_test, + test_size=test_size / (test_size + valid_size), + random_state=42, ) return ( x_train, @@ -157,8 +154,7 @@ def __init__( f"{self.model_name}_{self.tokenizer_name}_{self.data_name}" ) 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" ) if self.split_size is not None: self.train_size = self.split_size.get("train_size", 0.8) @@ -175,9 +171,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 classifier %s", data_name, self.model_name - ) + logging.info("Loading data: %s to classifier %s", data_name, self.model_name) self.data_name = data_name self.classifier_name = ( f"{self.model_name}_{self.tokenizer_name}_{self.data_name}" diff --git a/src/ai_detection/classifier/fusion_classifer.py b/src/ai_detection/classifier/fusion_classifer.py index b7483c0..d09dadf 100644 --- a/src/ai_detection/classifier/fusion_classifer.py +++ b/src/ai_detection/classifier/fusion_classifer.py @@ -4,13 +4,11 @@ import torch from lightning.pytorch.loggers import WandbLogger -from ai_detection.classifier.classify import TextClassifer, label_encode, split_data +from ai_detection.classifier.classify import (TextClassifer, label_encode, + split_data) from ai_detection.classifier.multi_feature_model.fusionBert import ( - CrossAttentionFeatureFusion, - FeatureFusionBertClassfier, - FeatureFusionDataModule, - HFeatureFusion, -) + CrossAttentionFeatureFusion, FeatureFusionBertClassfier, + FeatureFusionDataModule, HFeatureFusion) class PLClassifer(TextClassifer): diff --git a/src/ai_detection/classifier/huggingface_classifier.py b/src/ai_detection/classifier/huggingface_classifier.py index 1811072..8dfe56e 100644 --- a/src/ai_detection/classifier/huggingface_classifier.py +++ b/src/ai_detection/classifier/huggingface_classifier.py @@ -3,16 +3,11 @@ import datasets import evaluate import torch -from transformers import ( - DataCollatorWithPadding, - Trainer, - TrainingArguments, - BertForSequenceClassification, - BertTokenizer, -) +from transformers import (BertForSequenceClassification, BertTokenizer, + DataCollatorWithPadding, Trainer, TrainingArguments) - -from ai_detection.classifier.classify import TextClassifer, label_encode, split_data +from ai_detection.classifier.classify import (TextClassifer, label_encode, + split_data) from ai_detection.utils.everyai_path import MODEL_PATH @@ -34,9 +29,7 @@ def __init__( ) if self.model_name != "bert-base-uncased": raise ValueError(f"Model name {self.model_name} not supported yet") - self.model = BertForSequenceClassification.from_pretrained( - self.model_name - ) + self.model = BertForSequenceClassification.from_pretrained(self.model_name) self.tokenizer = BertTokenizer.from_pretrained( self.tokenizer_name, **self.tokenizer_config ) @@ -51,9 +44,7 @@ def __init__( def _tokenize(self, texts: list[str], labels: list[str]): self.label_encoder, tokenzied_labels = label_encode(labels) tokenzied_labels = torch.tensor(tokenzied_labels) - dataset = datasets.Dataset.from_dict( - {"text": texts, "label": tokenzied_labels} - ) + dataset = datasets.Dataset.from_dict({"text": texts, "label": tokenzied_labels}) def _tokenizer_fn(example): return self.tokenizer(example["text"], **self.tokenizer_config) @@ -77,12 +68,8 @@ def train(self): self.data.valid_indices, self.data.test_indices, ) = split_data(self.texts, self.labels) - self.train_dataset = self._tokenize( - self.data.x_train, self.data.y_train - ) - self.valid_dataset = self._tokenize( - self.data.x_valid, self.data.y_valid - ) + self.train_dataset = self._tokenize(self.data.x_train, self.data.y_train) + self.valid_dataset = self._tokenize(self.data.x_valid, self.data.y_valid) self.test_dataset = self._tokenize(self.data.x_test, self.data.y_test) train_args = TrainingArguments(**self.train_args) data_collator = DataCollatorWithPadding(tokenizer=self.tokenizer) @@ -98,9 +85,7 @@ def train(self): def test(self): trainer = Trainer(model=self.model) predictions = trainer.predict(self.test_dataset) - self.data.y_pred = torch.argmax( - torch.tensor(predictions.predictions), axis=1 - ) + self.data.y_pred = torch.argmax(torch.tensor(predictions.predictions), axis=1) self.data.y_test = self.label_encoder.transform(self.data.y_test) diff --git a/src/ai_detection/classifier/multi_feature_model/fusionBert.py b/src/ai_detection/classifier/multi_feature_model/fusionBert.py index d4de321..fed079f 100644 --- a/src/ai_detection/classifier/multi_feature_model/fusionBert.py +++ b/src/ai_detection/classifier/multi_feature_model/fusionBert.py @@ -24,18 +24,14 @@ def truncate_and_pad_single_sequence(seq, max_length): class FeatureFusionBertTokenizer: def __init__(self, feature_len, **kwargs): self.feature_len = feature_len - self.sentiment_max_length = kwargs.get( - "sentiment_max_length", feature_len - ) + self.sentiment_max_length = kwargs.get("sentiment_max_length", feature_len) self.nlp = spacy.load("en_core_web_sm") self.all_tags = self.nlp.get_pipe("tagger").labels self.sentiment_analyzer = SentimentIntensityAnalyzer() self.bert_model = BertForSequenceClassification.from_pretrained( "bert-base-uncased" ) - self.bert_tokenzier = BertTokenizer.from_pretrained( - "bert-base-uncased" - ) + self.bert_tokenzier = BertTokenizer.from_pretrained("bert-base-uncased") def semantic(self, text: str, **kwargs): tokenzied = self.bert_tokenzier( @@ -58,9 +54,9 @@ def analyze_word_level_sentiment(self, text: str) -> torch.tensor: if token.is_stop or token.is_punct: sentiment = 0.0 else: - sentiment = self.sentiment_analyzer.polarity_scores( - token.text - )["compound"] + sentiment = self.sentiment_analyzer.polarity_scores(token.text)[ + "compound" + ] word_sentiment.append(sentiment) sentiment = torch.tensor(word_sentiment, dtype=torch.float) return self._padding(sentiment) @@ -75,12 +71,8 @@ def pos_feature(self, text: str) -> torch.tensor: return self._padding(pos) def _padding(self, input_tensor): - input_tensor = truncate_and_pad_single_sequence( - input_tensor, self.feature_len - ) - attention_mask = torch.tensor( - [1] * len(input_tensor), dtype=torch.float - ) + input_tensor = truncate_and_pad_single_sequence(input_tensor, self.feature_len) + attention_mask = torch.tensor([1] * len(input_tensor), dtype=torch.float) return input_tensor.unsqueeze(0), attention_mask.unsqueeze(0) def __call__(self, text: str): @@ -112,9 +104,7 @@ def batch_encode_plus(self, batch_text: list[str]): # -> list: class HFeatureFusion(nn.Module): - def __init__( - self, feature_num, feature_len, output_dim, num_heads=4, dropout=0.1 - ): + def __init__(self, feature_num, feature_len, output_dim, num_heads=4, dropout=0.1): super().__init__() self.feature_num = feature_num self.feature_len = feature_len @@ -149,9 +139,7 @@ def forward(self, *features: torch.tensor): class CrossAttentionFeatureFusion(nn.Module): - def __init__( - self, feature_num, feature_len, output_dim, num_heads=4, dropout=0.1 - ): + def __init__(self, feature_num, feature_len, output_dim, num_heads=4, dropout=0.1): super().__init__() self.feature_len = feature_len self.num_heads = num_heads @@ -220,9 +208,7 @@ def __init__( def forward(self, *features): self.train() # Set all modules to train for feature in features: - feature["input_ids"] = F.normalize( - feature["input_ids"], p=2, dim=-1 - ) + feature["input_ids"] = F.normalize(feature["input_ids"], p=2, dim=-1) fused_features = self.fusion_module(*features) encoder_output = self.encoder(fused_features) pooler_output = self.pooler(encoder_output[0]) @@ -241,9 +227,7 @@ def test_step(self, batch, batch_idx): y_pred = torch.argmax(outputs, dim=-1) y_test = batch["labels"] acc = (y_pred == y_test).sum().float() / len(y_test) - f1 = f1_score( - y_test.cpu().numpy(), y_pred.cpu().numpy(), average="weighted" - ) + f1 = f1_score(y_test.cpu().numpy(), y_pred.cpu().numpy(), average="weighted") self.log("test_acc", acc) self.log("test_f1", f1) return {"test_acc": acc, "test_f1": f1} @@ -255,9 +239,7 @@ def test_epoch_end(self, outputs): self.log("avg_test_f1", avg_f1, prog_bar=True) def configure_optimizers(self): - optimizer = torch.optim.Adam( - self.parameters(), lr=self.lr, momentum=0.9 - ) + optimizer = torch.optim.Adam(self.parameters(), lr=self.lr, momentum=0.9) scheduler = get_scheduler( name="linear", optimizer=optimizer, @@ -318,8 +300,8 @@ def prepare_data(self): pass def setup(self, stage=None): - self.tokenizer: FeatureFusionBertTokenizer = ( - FeatureFusionBertTokenizer(feature_len=768) + self.tokenizer: FeatureFusionBertTokenizer = FeatureFusionBertTokenizer( + feature_len=768 ) if not isinstance(self.tokenizer, FeatureFusionBertTokenizer): raise TypeError( diff --git a/src/ai_detection/classifier/multi_feature_model/fusionBert_backup.py b/src/ai_detection/classifier/multi_feature_model/fusionBert_backup.py index 80b4e77..74fcec3 100644 --- a/src/ai_detection/classifier/multi_feature_model/fusionBert_backup.py +++ b/src/ai_detection/classifier/multi_feature_model/fusionBert_backup.py @@ -24,18 +24,14 @@ def truncate_and_pad_single_sequence(seq, max_length): class FeatureFusionBertTokenizer: def __init__(self, feature_len, **kwargs): self.feature_len = feature_len - self.sentiment_max_length = kwargs.get( - "sentiment_max_length", feature_len - ) + self.sentiment_max_length = kwargs.get("sentiment_max_length", feature_len) self.nlp = spacy.load("en_core_web_sm") self.all_tags = self.nlp.get_pipe("tagger").labels self.sentiment_analyzer = SentimentIntensityAnalyzer() self.bert_model = BertForSequenceClassification.from_pretrained( "bert-base-uncased" ) - self.bert_tokenzier = BertTokenizer.from_pretrained( - "bert-base-uncased" - ) + self.bert_tokenzier = BertTokenizer.from_pretrained("bert-base-uncased") def semantic(self, text: str, **kwargs): tokenzied = self.bert_tokenzier( @@ -58,9 +54,9 @@ def analyze_word_level_sentiment(self, text: str) -> torch.tensor: if token.is_stop or token.is_punct: sentiment = 0.0 else: - sentiment = self.sentiment_analyzer.polarity_scores( - token.text - )["compound"] + sentiment = self.sentiment_analyzer.polarity_scores(token.text)[ + "compound" + ] word_sentiment.append(sentiment) sentiment = torch.tensor(word_sentiment, dtype=torch.float) return self._padding(sentiment) @@ -75,12 +71,8 @@ def pos_feature(self, text: str) -> torch.tensor: return self._padding(pos) def _padding(self, input_tensor): - input_tensor = truncate_and_pad_single_sequence( - input_tensor, self.feature_len - ) - attention_mask = torch.tensor( - [1] * len(input_tensor), dtype=torch.float - ) + input_tensor = truncate_and_pad_single_sequence(input_tensor, self.feature_len) + attention_mask = torch.tensor([1] * len(input_tensor), dtype=torch.float) return input_tensor.unsqueeze(0), attention_mask.unsqueeze(0) def __call__(self, text: str): @@ -111,9 +103,7 @@ def batch_encode_plus(self, batch_text: list[str], **kwargs): class HFeatureFusion(nn.Module): - def __init__( - self, feature_num, feature_len, output_dim, num_heads=4, dropout=0.1 - ): + def __init__(self, feature_num, feature_len, output_dim, num_heads=4, dropout=0.1): super().__init__() self.feature_num = feature_num self.feature_len = feature_len @@ -151,9 +141,7 @@ def forward(self, *features: torch.tensor): class CrossAttentionFeatureFusion(nn.Module): - def __init__( - self, feature_num, feature_len, output_dim, num_heads=4, dropout=0.1 - ): + def __init__(self, feature_num, feature_len, output_dim, num_heads=4, dropout=0.1): super().__init__() self.feature_len = feature_len self.num_heads = num_heads @@ -255,19 +243,17 @@ def __len__(self): for epoch in range(num_epochs): for step, batch in enumerate(train_dataloader): for feature in ["semantic", "pos", "sentiment"]: - batch[feature]["input_ids"] = batch[feature]["input_ids"].to( + batch[feature]["input_ids"] = batch[feature]["input_ids"].to(device) + batch[feature]["attention_mask"] = batch[feature]["attention_mask"].to( device ) - batch[feature]["attention_mask"] = batch[feature][ - "attention_mask" - ].to(device) optimizer.zero_grad() - outputs = model( - batch["semantic"], batch["pos"], batch["sentiment"] - ) + outputs = model(batch["semantic"], batch["pos"], batch["sentiment"]) loss = F.cross_entropy(outputs, batch["labels"].to(device)) loss.backward() optimizer.step() lr_scheduler.step() if step % 100 == 0: - print(f"Epoch {epoch + 1}, step {step + 1} completed with loss: {loss.item()}") + print( + f"Epoch {epoch + 1}, step {step + 1} completed with loss: {loss.item()}" + ) diff --git a/src/ai_detection/classifier/pytorch_classier.py b/src/ai_detection/classifier/pytorch_classier.py index 6fe1f7e..2ad81bf 100644 --- a/src/ai_detection/classifier/pytorch_classier.py +++ b/src/ai_detection/classifier/pytorch_classier.py @@ -6,7 +6,8 @@ from transformers import (AutoTokenizer, DataCollatorWithPadding, Trainer, TrainingArguments) -from ai_detection.classifier.classify import TextClassifer, label_encode, split_data +from ai_detection.classifier.classify import (TextClassifer, label_encode, + split_data) from ai_detection.classifier.multi_feature_model.fusionBert_backup import ( FeatureFusionBertClassfier, FeatureFusionBertTokenizer) from ai_detection.utils.everyai_path import MODEL_PATH @@ -28,13 +29,13 @@ def __init__( language=language, **classfiy_config, ) - tokenzier_dict = { - "fusion_bert": FeatureFusionBertTokenizer(AutoTokenizer("bert-base-uncased")) + tokenzier_dict = { + "fusion_bert": FeatureFusionBertTokenizer( + AutoTokenizer("bert-base-uncased") + ) } self.tokenizer = tokenzier_dict[self.tokenizer_name] - model_dict = { - "fusion_bert": FeatureFusionBertClassfier(feature_num=3) - } + model_dict = {"fusion_bert": FeatureFusionBertClassfier(feature_num=3)} self.model = model_dict[self.model_name] self.label_encoder = None self.train_dataset, self.valid_dataset, self.test_dataset = ( @@ -47,9 +48,7 @@ def __init__( def _tokenize(self, texts: list[str], labels: list[str]): self.label_encoder, tokenzied_labels = label_encode(labels) tokenzied_labels = torch.tensor(tokenzied_labels) - dataset = datasets.Dataset.from_dict( - {"text": texts, "label": tokenzied_labels} - ) + dataset = datasets.Dataset.from_dict({"text": texts, "label": tokenzied_labels}) def _tokenizer_fn(example): return self.tokenizer(example["text"], **self.tokenizer_config) @@ -73,12 +72,8 @@ def train(self): self.data.valid_indices, self.data.test_indices, ) = split_data(self.texts, self.labels) - self.train_dataset = self._tokenize( - self.data.x_train, self.data.y_train - ) - self.valid_dataset = self._tokenize( - self.data.x_valid, self.data.y_valid - ) + self.train_dataset = self._tokenize(self.data.x_train, self.data.y_train) + self.valid_dataset = self._tokenize(self.data.x_valid, self.data.y_valid) self.test_dataset = self._tokenize(self.data.x_test, self.data.y_test) train_args = TrainingArguments(**self.train_args) data_collator = DataCollatorWithPadding(tokenizer=self.tokenizer) @@ -94,15 +89,11 @@ def train(self): def test(self): trainer = Trainer(model=self.model) predictions = trainer.predict(self.test_dataset) - self.data.y_pred = torch.argmax( - torch.tensor(predictions.predictions), axis=1 - ) + self.data.y_pred = torch.argmax(torch.tensor(predictions.predictions), axis=1) self.data.y_test = self.label_encoder.transform(self.data.y_test) def show_score(self): metric = evaluate.load("accuracy") - metric.compute( - predictions=self.data.y_pred, references=self.data.y_test - ) + metric.compute(predictions=self.data.y_pred, references=self.data.y_test) logging.info("Accuracy: %s", metric) diff --git a/src/ai_detection/classifier/sklearn_classifier.py b/src/ai_detection/classifier/sklearn_classifier.py index c828ea6..74d109e 100644 --- a/src/ai_detection/classifier/sklearn_classifier.py +++ b/src/ai_detection/classifier/sklearn_classifier.py @@ -20,7 +20,8 @@ from sklearn.svm import SVC from sklearn.tree import DecisionTreeClassifier -from ai_detection.classifier.classify import TextClassifer, label_encode, split_data +from ai_detection.classifier.classify import (TextClassifer, label_encode, + split_data) from ai_detection.utils.everyai_path import RESULT_PATH @@ -141,9 +142,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]) @@ -234,7 +233,11 @@ def predict(self, x): return self.model.predict(x) def show_score(self): - output_path = RESULT_PATH/ "classfiy_result"/ f"{self.model_name}_{self.tokenizer_name}_{self.data_name}" + output_path = ( + RESULT_PATH + / "classfiy_result" + / f"{self.model_name}_{self.tokenizer_name}_{self.data_name}" + ) self.score = evaluate_classification_model( self.data.y_test, self.data.y_pred, diff --git a/src/ai_detection/data_loader/data_load.py b/src/ai_detection/data_loader/data_load.py index d3548a6..6d286d7 100644 --- a/src/ai_detection/data_loader/data_load.py +++ b/src/ai_detection/data_loader/data_load.py @@ -65,9 +65,7 @@ def load_data( if max_count is not None and loaded_data is not None: loaded_data = loaded_data.head(max_count) else: - logging.info( - "Max count is None and all the records will be loaded" - ) + logging.info("Max count is None and all the records will be loaded") loaded_data.rename( columns={ self.question_column: "question", @@ -78,9 +76,7 @@ def load_data( if return_type == "pandas": result = loaded_data elif return_type == "list": - result = loaded_data[["question", "answer"]].to_dict( - orient="records" - ) + result = loaded_data[["question", "answer"]].to_dict(orient="records") else: logging.error("Invalid return type") return result @@ -90,4 +86,4 @@ def apply_filter(self, orginal_data: pd.DataFrame) -> pd.DataFrame: orginal_data if self.data_filter is None else orginal_data[orginal_data.apply(self.data_filter, axis=1)] - ) \ No newline at end of file + ) diff --git a/src/ai_detection/data_loader/data_process.py b/src/ai_detection/data_loader/data_process.py index ed581cf..cb78b94 100644 --- a/src/ai_detection/data_loader/data_process.py +++ b/src/ai_detection/data_loader/data_process.py @@ -5,7 +5,8 @@ import jieba -from ai_detection.utils.everyai_path import EN_STOP_WORD_PATH, ZH_STOP_WORD_PATH +from ai_detection.utils.everyai_path import (EN_STOP_WORD_PATH, + ZH_STOP_WORD_PATH) def remove_punctuation(text: str) -> str: diff --git a/src/ai_detection/data_loader/everyai_dataset.py b/src/ai_detection/data_loader/everyai_dataset.py index ef3ec9d..bd99598 100644 --- a/src/ai_detection/data_loader/everyai_dataset.py +++ b/src/ai_detection/data_loader/everyai_dataset.py @@ -1,12 +1,13 @@ import logging -from pathlib import Path import re +from pathlib import Path import pandas as pd import pymongo import pymongo.database -from ai_detection.data_loader.data_process import split_remove_stopwords_punctuation +from ai_detection.data_loader.data_process import \ + split_remove_stopwords_punctuation from ai_detection.data_loader.mongo_connection import get_mongo_connection from ai_detection.utils.everyai_path import DATA_PATH, MONGO_CONFIG_PATH from ai_detection.utils.load_config import get_config @@ -32,9 +33,7 @@ def __init__( if datas is not None: self.datas: pd.DataFrame = datas else: - self.datas: pd.DataFrame = pd.DataFrame( - columns=["question", "human"] - ) + self.datas: pd.DataFrame = pd.DataFrame(columns=["question", "human"]) if ai_list is not None: for ai_name in ai_list: self.datas[ai_name] = None @@ -87,33 +86,25 @@ def insert_ai_response(self, question, ai_name: str, ai_response: str): else: logging.info("AI %s exists in the dataset", ai_name) question_exists = not self.datas[ - (self.datas["question"] == question) - & (self.datas[ai_name] == ai_name) + (self.datas["question"] == question) & (self.datas[ai_name] == ai_name) ].empty if question_exists: self._update_new_row(question, ai_name, ai_response) else: - self.datas.loc[self.datas["question"] == question, ai_name] = ( - ai_response - ) + self.datas.loc[self.datas["question"] == question, ai_name] = ai_response def insert_human_response(self, question, human_response: str): if self.datas[self.datas["question"] == question].empty: self._update_new_row(question, "human", human_response) else: - self.datas.loc[self.datas["question"] == question, "human"] = ( - human_response - ) + self.datas.loc[self.datas["question"] == question, "human"] = human_response - def record_exist( - self, question: str, label: str, file_type: str = "mongodb" - ): + def record_exist(self, question: str, label: str, file_type: str = "mongodb"): if file_type != "mongodb": raise ValueError("Unsupported file type: %s", file_type) query = {"question": question, label: {"$exists": True}} record = self.mongo_collection.find_one(query) return False if record is None else isinstance(record[label], str) - def upsert2mongo(self, question: str, label: str, answer: str): query = {"question": question} @@ -128,9 +119,7 @@ def _update_new_row(self, question, arg1, arg2): def output_question(self): # -> Iterator: return iter(self.datas["question"]) - def _save2mongodb( - self, database: pymongo.database.Database, insert_mode="insert" - ): + def _save2mongodb(self, database: pymongo.database.Database, insert_mode="insert"): self.datas = self.datas.dropna() logging.info("Saving dataset to mongodb: %s", database) collection = database[self.data_name] @@ -160,9 +149,7 @@ def _read_from_mongodb(self, database: pymongo.database.Database): data = data.drop(columns=["timestamp"]) return data - def read( - self, path_or_database: str | Path = None, file_format: str = "csv" - ): + def read(self, path_or_database: str | Path = None, file_format: str = "csv"): if file_format == "mongodb": if path_or_database is None: path_or_database = _initialize_mongo_connection() @@ -171,17 +158,13 @@ def read( loaded_data = self._read_from_mongodb(path_or_database) else: if path_or_database is None: - path_or_database = ( - DATA_PATH / f"{self.data_name}.{file_format}" - ) + path_or_database = DATA_PATH / f"{self.data_name}.{file_format}" logging.info("Load dataset from %s", path_or_database) if isinstance(path_or_database, str): path_or_database = Path(path_or_database) if path_or_database.suffix != f".{file_format}": logging.warning("Change file format to %s", file_format) - path_or_database = path_or_database.with_suffix( - f".{file_format}" - ) + path_or_database = path_or_database.with_suffix(f".{file_format}") match path_or_database.suffix: case ".csv": loaded_data = pd.read_csv(path_or_database) @@ -204,9 +187,7 @@ def read( self.datas.dropna(inplace=True) self.ai_list = list(set(self.datas.columns) - {"question", "human"}) - def save( - self, path_or_database: str | Path = None, file_format: str = "csv" - ): + def save(self, path_or_database: str | Path = None, file_format: str = "csv"): self.datas = self.datas.dropna() if file_format == "mongodb": if path_or_database is None: @@ -214,17 +195,13 @@ def save( self._save2mongodb(path_or_database) else: if path_or_database is None: - path_or_database = ( - DATA_PATH / f"{self.data_name}.{file_format}" - ) + path_or_database = DATA_PATH / f"{self.data_name}.{file_format}" logging.info("Save dataset to %s", path_or_database) if isinstance(path_or_database, str): path_or_database = Path(path_or_database) if path_or_database.suffix != f".{file_format}": logging.warning("Change file format to %s", file_format) - path_or_database = path_or_database.with_suffix( - f".{file_format}" - ) + path_or_database = path_or_database.with_suffix(f".{file_format}") match path_or_database.suffix: case ".csv": self.datas.to_csv(path_or_database, index=False) @@ -235,8 +212,8 @@ def save( case _: logging.error("Invalid format: %s", file_format) return path_or_database - - def export4train(self,path:str|Path) -> pd.DataFrame: + + def export4train(self, path: str | Path) -> pd.DataFrame: """Export data for training. Combines the texts and labels into a DataFrame suitable for training. @@ -248,4 +225,3 @@ def export4train(self,path:str|Path) -> pd.DataFrame: data = pd.DataFrame({"text": texts, "label": labels}) if isinstance(path, str): path = Path(path) - diff --git a/src/ai_detection/data_loader/mongo_connection.py b/src/ai_detection/data_loader/mongo_connection.py index ebfdbe5..05ac6af 100644 --- a/src/ai_detection/data_loader/mongo_connection.py +++ b/src/ai_detection/data_loader/mongo_connection.py @@ -5,7 +5,8 @@ from pymongo import MongoClient -def get_mongo_connection(connection_string:str, database_name:str):# -> Database: +# -> Database: +def get_mongo_connection(connection_string: str, database_name: str): if connection_string is None or not connection_string: logging.info("Use the connection string in environment variable") connection_string = os.getenv("MONGO_CONNECTION_STRING") diff --git a/src/ai_detection/explanation/explain.py b/src/ai_detection/explanation/explain.py index dce9573..27a0c7b 100644 --- a/src/ai_detection/explanation/explain.py +++ b/src/ai_detection/explanation/explain.py @@ -36,7 +36,9 @@ def explain(self, output_path: Path = None): exp.save_to_file(output_path / f"text{i}.html") logging.info("Lime explanation saved to %s", output_path) else: - logging.error("Lime explanation not supported for %s", self.classifier.classifier_type) + logging.error( + "Lime explanation not supported for %s", self.classifier.classifier_type + ) class ShapExplanation(Explanation): @@ -57,4 +59,6 @@ def explain(self, output_path: Path = None): shap.plots.text(shap_values, output_path / f"text{i}") logging.info("Shap explanation saved to %s", output_path) else: - logging.error("Shap explanation not supported for %s", self.classifier.classifier_type) \ No newline at end of file + logging.error( + "Shap explanation not supported for %s", self.classifier.classifier_type + ) diff --git a/src/ai_detection/generator/generate.py b/src/ai_detection/generator/generate.py index 61f33e3..5df4748 100644 --- a/src/ai_detection/generator/generate.py +++ b/src/ai_detection/generator/generate.py @@ -2,8 +2,8 @@ from collections.abc import Callable from pathlib import Path -from ai_detection.generator.openai_generate import openai_generate from ai_detection.generator.huggingface_generate import glm4, llama, qwen2_5 +from ai_detection.generator.openai_generate import openai_generate from ai_detection.utils.everyai_path import GENERATE_CONFIG_PATH from ai_detection.utils.load_args import set_attrs_2class from ai_detection.utils.load_config import get_config @@ -37,9 +37,7 @@ def __init__(self, config: dict, formatter: Callable[[str], str] = None): ] default_params = ["model_name", "base_url", "api_key"] else: - raise ValueError( - "Unsupported generator type: %s", self.generator_type - ) + raise ValueError("Unsupported generator type: %s", self.generator_type) self.formatter = formatter set_attrs_2class(self, config, allowed_keys, default_params) diff --git a/src/ai_detection/generator/huggingface_generate.py b/src/ai_detection/generator/huggingface_generate.py index cb19dbb..f0491a6 100644 --- a/src/ai_detection/generator/huggingface_generate.py +++ b/src/ai_detection/generator/huggingface_generate.py @@ -2,8 +2,8 @@ from pathlib import Path import torch -from transformers import AutoModelForCausalLM, AutoTokenizer import transformers +from transformers import AutoModelForCausalLM, AutoTokenizer def glm4( @@ -38,7 +38,7 @@ def glm4( return_tensors="pt", add_generation_prompt=True, return_dict=True, - template="default" # Add this line to specify the template + template="default", # Add this line to specify the template ).to(model.device) input_len = inputs["input_ids"].shape[1] @@ -81,13 +81,15 @@ def qwen2_5( }, {"role": "user", "content": user_input}, ] - + text = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) else: text = user_input - logging.warning("No chat template found in tokenizer. Using user input as text.") + logging.warning( + "No chat template found in tokenizer. Using user input as text." + ) model_inputs = tokenizer([text], return_tensors="pt").to(model.device) generated_ids = model.generate(**model_inputs, **gen_kwargs) diff --git a/src/ai_detection/generator/openai_generate.py b/src/ai_detection/generator/openai_generate.py index c7ca7c9..d6ecc35 100644 --- a/src/ai_detection/generator/openai_generate.py +++ b/src/ai_detection/generator/openai_generate.py @@ -1,6 +1,7 @@ import logging from openai import OpenAI + from ai_detection.utils.proxy import TempProxy @@ -17,9 +18,7 @@ def openai_generate( logging.info("Proxy started: %s", proxy) client = OpenAI(api_key=api_key, base_url=base_url) messages = [{"role": "user", "content": user_input}] - result = client.chat.completions.create( - messages=messages, model=model_name - ) + result = client.chat.completions.create(messages=messages, model=model_name) if proxy is not None: temp_proxy.reset_proxy() logging.info("Proxy reset") diff --git a/src/ai_detection/main.py b/src/ai_detection/main.py index 8b31aef..2192255 100644 --- a/src/ai_detection/main.py +++ b/src/ai_detection/main.py @@ -8,15 +8,16 @@ from ai_detection.classifier.huggingface_classifier import HuggingfaceClassifer from ai_detection.classifier.sklearn_classifier import SklearnClassifer from ai_detection.data_loader.data_load import Data_loader -from ai_detection.data_loader.data_process import split_remove_stopwords_punctuation +from ai_detection.data_loader.data_process import \ + split_remove_stopwords_punctuation from ai_detection.data_loader.everyai_dataset import EveryaiDataset from ai_detection.explanation.explain import LimeExplanation, ShapExplanation from ai_detection.generator.generate import Generator from ai_detection.topic.my_bertopic import create_topic from ai_detection.utils.everyai_path import (BERT_TOPIC_CONFIG_PATH, - CLASSIFY_CONFIG_PATH, - DATA_LOAD_CONFIG_PATH, DATA_PATH, - FIG_PATH, GENERATE_CONFIG_PATH) + CLASSIFY_CONFIG_PATH, + DATA_LOAD_CONFIG_PATH, DATA_PATH, + FIG_PATH, GENERATE_CONFIG_PATH) from ai_detection.utils.load_config import get_config @@ -30,9 +31,7 @@ def generate(): data_loader = Data_loader( **data_config, ) - max_count = ( - data_config["max_count"] if "max_count" in data_config else None - ) + max_count = data_config["max_count"] if "max_count" in data_config else None qa_datas = data_loader.load_data(max_count=max_count) everyai_dataset = EveryaiDataset( datas=pd.DataFrame(qa_datas), @@ -45,9 +44,7 @@ def generate(): for generate_config in generate_list_configs["generate_list"]: logging.info("Generate config: %s", generate_config) generator = Generator(config=generate_config) - 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)): if everyai_dataset.record_exist( data["question"], generate_config["model_name"] ): diff --git a/src/ai_detection/test.py b/src/ai_detection/test.py index f1b54d7..8cd94aa 100644 --- a/src/ai_detection/test.py +++ b/src/ai_detection/test.py @@ -1,2 +1,3 @@ from modelscope import snapshot_download -model_dir = snapshot_download('YIRONGCHEN/SoulChat2.0-Llama-3.1-8B') + +model_dir = snapshot_download("YIRONGCHEN/SoulChat2.0-Llama-3.1-8B") diff --git a/src/ai_detection/topic/my_bertopic.py b/src/ai_detection/topic/my_bertopic.py index c2e0727..f3d1c72 100644 --- a/src/ai_detection/topic/my_bertopic.py +++ b/src/ai_detection/topic/my_bertopic.py @@ -10,8 +10,9 @@ def create_topic( output_folder: Union[str, Path], embedding_model=None, topic_config: dict = None, -) : +): from bertopic import BERTopic + if topic_config is None: topic_config = {} topic_model = BERTopic(embedding_model=embedding_model, min_topic_size=5) diff --git a/src/ai_detection/utils/load_args.py b/src/ai_detection/utils/load_args.py index e6751ad..2d4231a 100644 --- a/src/ai_detection/utils/load_args.py +++ b/src/ai_detection/utils/load_args.py @@ -11,7 +11,9 @@ ] -def set_attrs_2class(self, classify_config:dict, allowed_keys:list, default_keys:list): +def set_attrs_2class( + self, classify_config: dict, allowed_keys: list, default_keys: list +): """Set attributes for classification. Sets attributes based on the provided configuration, default values, and necessary keys. @@ -31,4 +33,3 @@ def set_attrs_2class(self, classify_config:dict, allowed_keys:list, default_keys if key not in classify_config: setattr(self, key, None) logging.warning("Default key not provided: %s", key) - diff --git a/src/ai_detection/utils/load_config.py b/src/ai_detection/utils/load_config.py index 85e1b64..c354d15 100644 --- a/src/ai_detection/utils/load_config.py +++ b/src/ai_detection/utils/load_config.py @@ -2,13 +2,11 @@ import yaml -from ai_detection.utils.everyai_path import ( - BERT_TOPIC_CONFIG_PATH, - CLASSIFY_CONFIG_PATH, - DATA_LOAD_CONFIG_PATH, - GENERATE_CONFIG_PATH, - MONGO_CONFIG_PATH, -) +from ai_detection.utils.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):