diff --git a/data/classification/.gitignore b/data/classification/.gitignore
index 60ba70084bebf52a8521c364d3a4b019028fe084..e69587204a79358e2809c210d36b6519d6be3b81 100644
--- a/data/classification/.gitignore
+++ b/data/classification/.gitignore
@@ -1 +1,3 @@
/enron_spam
+/wiki_pl
+/20_news
diff --git a/data/datasets/.gitignore b/data/datasets/.gitignore
index af871dfd6ce1c3aab7e8d1a405df6390acab6f65..43bd163e72f8f4c3216a67dd67aaa506107fbda5 100644
--- a/data/datasets/.gitignore
+++ b/data/datasets/.gitignore
@@ -1,2 +1,4 @@
/enron_spam
+/20_news
/poleval
+/wiki_pl
diff --git a/data/datasets/20_news.dvc b/data/datasets/20_news.dvc
new file mode 100644
index 0000000000000000000000000000000000000000..00b5cf40552d824c249c9e692753ce5dbdf3b4d5
--- /dev/null
+++ b/data/datasets/20_news.dvc
@@ -0,0 +1,5 @@
+outs:
+- md5: 999207f1c2c123c9943397b47f2c3b3a.dir
+ size: 23460358
+ nfiles: 3
+ path: 20_news
diff --git a/data/datasets/wiki_pl.dvc b/data/datasets/wiki_pl.dvc
new file mode 100644
index 0000000000000000000000000000000000000000..f0f2afeb3fb36b5b1a88e081bc18151e9f3500dd
--- /dev/null
+++ b/data/datasets/wiki_pl.dvc
@@ -0,0 +1,5 @@
+outs:
+- md5: abcbccb3e352ed623cace1b95078bd63.dir
+ size: 29115538
+ nfiles: 3
+ path: wiki_pl
diff --git a/data/models/.gitignore b/data/models/.gitignore
index 60ba70084bebf52a8521c364d3a4b019028fe084..ea22867615bba98d219c12d7f14467d051a33e80 100644
--- a/data/models/.gitignore
+++ b/data/models/.gitignore
@@ -1 +1,3 @@
/enron_spam
+/20_news
+/wiki_pl
diff --git a/data/models/20_news.dvc b/data/models/20_news.dvc
new file mode 100644
index 0000000000000000000000000000000000000000..d667d5706620ffdbfc2e6148aefc3a781540abf4
--- /dev/null
+++ b/data/models/20_news.dvc
@@ -0,0 +1,5 @@
+outs:
+- md5: 43d68a67ecb8149bd6bf50db9767cb64.dir
+ size: 439008808
+ nfiles: 6
+ path: 20_news
diff --git a/data/models/wiki_pl.dvc b/data/models/wiki_pl.dvc
new file mode 100644
index 0000000000000000000000000000000000000000..fdf58d54d28455296247165dff7d827def75296c
--- /dev/null
+++ b/data/models/wiki_pl.dvc
@@ -0,0 +1,5 @@
+outs:
+- md5: fd453042628fb09c080ef05d34a32cce.dir
+ size: 501711136
+ nfiles: 7
+ path: wiki_pl
diff --git a/experiments/scripts/classify.py b/experiments/scripts/classify.py
index 9639d298904c1a2815f0b34f8cbb6894df6c8527..ab34bd70e815f74effd4044efa041f3ebb5249d6 100644
--- a/experiments/scripts/classify.py
+++ b/experiments/scripts/classify.py
@@ -3,6 +3,7 @@ from pathlib import Path
import click
import pandas as pd
+import torch
from sklearn.metrics import classification_report
from text_attacks.utils import get_classify_function
@@ -27,6 +28,7 @@ def main(
output_dir.mkdir(parents=True, exist_ok=True)
classify = get_classify_function(
dataset_name=dataset_name,
+ device="cuda" if torch.cuda.is_available() else "cpu"
)
test = pd.read_json(f"data/preprocessed/{dataset_name}/test.jsonl", lines=True)
test_x = test["text"].tolist()
diff --git a/requirements.txt b/requirements.txt
index fec55bda346006486c107831bf7b436564332a9e..66b509aeeee23df8984167e7e05aa922e01fbfa2 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,8 +2,8 @@ datasets
transformers
click
scikit-learn
-dvc[s3]
-shap
+dvc[s3]==2.46.0
+shap==0.41.0
lpmn_client_biz
--find-links https://download.pytorch.org/whl/torch_stable.html
diff --git a/text_attacks/models/20_news.py b/text_attacks/models/20_news.py
new file mode 100644
index 0000000000000000000000000000000000000000..53712fa403a55b03a12aaf6962cdfcd6f4c503ce
--- /dev/null
+++ b/text_attacks/models/20_news.py
@@ -0,0 +1,42 @@
+"""Classification model for enron_spam"""
+import os
+
+import torch
+from tqdm import tqdm
+
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+
+def get_model_and_tokenizer():
+ model_path = "./data/models/20_news"
+ tokenizer = AutoTokenizer.from_pretrained(model_path)
+ model = AutoModelForSequenceClassification.from_pretrained(model_path)
+ return model, tokenizer
+
+
+def get_classify_function(device="cpu"):
+ model, tokenizer = get_model_and_tokenizer()
+ model.eval()
+ model = model.to(device)
+
+ def fun(texts):
+ logits = list()
+ i = 0
+ for chunk in tqdm(
+ [texts[pos:pos + 256] for pos in range(0, len(texts), 256)]
+ ):
+ encoded_inputs = tokenizer(
+ chunk,
+ return_tensors="pt",
+ padding=True,
+ truncation=True,
+ max_length=512
+ ).to(device)
+ with torch.no_grad():
+ logits.append(model(**encoded_inputs).logits.cpu())
+ logits = torch.cat(logits, dim=0)
+ pred_y = torch.argmax(logits, dim=1).tolist()
+ pred_y = [model.config.id2label[p] for p in pred_y]
+ return pred_y
+
+ return fun
diff --git a/text_attacks/models/enron_spam.py b/text_attacks/models/enron_spam.py
index 063a52a0c6cb7f09a804e00f19fc90d69944aa0e..9a1946d83a2ab8ad886d66ee8303cb283d74884b 100644
--- a/text_attacks/models/enron_spam.py
+++ b/text_attacks/models/enron_spam.py
@@ -2,12 +2,13 @@
import os
import torch
+from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForSequenceClassification
def get_model_and_tokenizer():
- model_path = "data/models/endron_spam"
+ model_path = "./data/models/endron_spam"
if not os.path.exists(model_path):
model_path = "mrm8488/bert-tiny-finetuned-enron-spam-detection"
tokenizer = AutoTokenizer.from_pretrained(model_path)
@@ -16,18 +17,27 @@ def get_model_and_tokenizer():
return model, tokenizer
-def get_classify_function():
+def get_classify_function(device="cpu"):
model, tokenizer = get_model_and_tokenizer()
+ model.eval()
+ model = model.to(device)
def fun(texts):
- encoded_inputs = tokenizer(
- texts,
- return_tensors="pt",
- padding=True,
- truncation=True,
- max_length=512
- )
- logits = model(**encoded_inputs).logits
+ logits = list()
+ i = 0
+ for chunk in tqdm(
+ [texts[pos:pos + 256] for pos in range(0, len(texts), 256)]
+ ):
+ encoded_inputs = tokenizer(
+ chunk,
+ return_tensors="pt",
+ padding=True,
+ truncation=True,
+ max_length=512
+ ).to(device)
+ with torch.no_grad():
+ logits.append(model(**encoded_inputs).logits.cpu())
+ logits = torch.cat(logits, dim=0)
pred_y = torch.argmax(logits, dim=1).tolist()
pred_y = [model.config.id2label[p] for p in pred_y]
return pred_y
diff --git a/text_attacks/models/wiki_pl.py b/text_attacks/models/wiki_pl.py
new file mode 100644
index 0000000000000000000000000000000000000000..1ad153955d6e5acfc89f4f922465fb624c1ecf5d
--- /dev/null
+++ b/text_attacks/models/wiki_pl.py
@@ -0,0 +1,42 @@
+"""Classification model for enron_spam"""
+import os
+
+import torch
+from tqdm import tqdm
+
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+
+def get_model_and_tokenizer():
+ model_path = "./data/models/wiki_pl"
+ tokenizer = AutoTokenizer.from_pretrained(model_path)
+ model = AutoModelForSequenceClassification.from_pretrained(model_path)
+ return model, tokenizer
+
+
+def get_classify_function(device="cpu"):
+ model, tokenizer = get_model_and_tokenizer()
+ model.eval()
+ model = model.to(device)
+
+ def fun(texts):
+ logits = list()
+ i = 0
+ for chunk in tqdm(
+ [texts[pos:pos + 256] for pos in range(0, len(texts), 256)]
+ ):
+ encoded_inputs = tokenizer(
+ chunk,
+ return_tensors="pt",
+ padding=True,
+ truncation=True,
+ max_length=512
+ ).to(device)
+ with torch.no_grad():
+ logits.append(model(**encoded_inputs).logits.cpu())
+ logits = torch.cat(logits, dim=0)
+ pred_y = torch.argmax(logits, dim=1).tolist()
+ pred_y = [model.config.id2label[p] for p in pred_y]
+ return pred_y
+
+ return fun
diff --git a/text_attacks/utils.py b/text_attacks/utils.py
index e47d5209c6c84c6b31d6836aef75051a6c66b57f..6a0588292c9542cc5b2d8bb5bd6a1437150276d0 100644
--- a/text_attacks/utils.py
+++ b/text_attacks/utils.py
@@ -11,13 +11,13 @@ def get_model_and_tokenizer(dataset_name):
return fun()
-def get_classify_function(dataset_name):
+def get_classify_function(dataset_name, device="cpu"):
"""Return get_model_and_tokenizer for a specific dataset."""
fun = getattr(
importlib.import_module(f"text_attacks.models.{dataset_name}"),
"get_classify_function",
)
- return fun()
+ return fun(device=device)
def download_dataset(dataset_name):