Format

Signed-off-by: András Schmelczer <andras@schmelczer.dev>

example/main.py

@@ -1,8 +1,8 @@
 import json
 
-from good_ai import LargeFile, process_batch
 from predict_domain import predict_domain
 
+from good_ai import LargeFile, process_batch
 
 if __name__ == "__main__":
     with open("data/s2-corpus-1583.json") as f:

example/predict_domain.py

@@ -2,13 +2,14 @@ import re
 from typing import Dict, Iterable, List
 
 from config import model_key
-from preprocess import preprocess
 from models import DomainPrediction
+from preprocess import preprocess
 from sklearn.pipeline import Pipeline
 
 from good_ai import use_model
 from good_ai.utilities.clean import clean
 
+
 @use_model(model_key, version="latest")
 def predict_domain(
     text: str, model: Pipeline, cut_off_probability: float = 0.2
@@ -19,12 +20,11 @@ def predict_domain(
 
     feature_names = model.named_steps["vectorizer"].get_feature_names_out()
 
-    token_mapping = {
+    token_mapping = {preprocess(original): original for original in text.split(" ")}
-        preprocess(original): original
-        for original in text.split(" ")
-    }
 
-    features = model.named_steps["vectorizer"].transform([" ".join(token_mapping.keys())])
+    features = model.named_steps["vectorizer"].transform(
+        [" ".join(token_mapping.keys())]
+    )
     prediction = model.named_steps["classifier"].predict_proba(features)[0]
     best_classes = sorted(enumerate(prediction), key=lambda v: v[1], reverse=True)
 

example/preprocess.py

@@ -1,6 +1,5 @@
 import re
 
-from good_ai.utilities.clean import clean
 from good_ai.utilities.lemmatize_text import lemmatize_text
 
 

example/train.ipynb

@@ -119,7 +119,9 @@
     "            content = json.load(f)\n",
     "\n",
     "        result = {\n",
-    "            preprocess(clean(f'{c[\"title\"]} {c[\"abstract\"]}'), convert_to_ascii=True): c[\"domain\"]\n",
+    "            preprocess(\n",
+    "                clean(f'{c[\"title\"]} {c[\"abstract\"]}'), convert_to_ascii=True\n",
+    "            ): c[\"domain\"]\n",
     "            for c in content\n",
     "            if (\n",
     "                c[\"domain\"]\n",
@@ -242,7 +244,7 @@
    "source": [
     "classifier = Pipeline(\n",
     "    steps=[\n",
-    "        (\"vectorizer\", TfidfVectorizer(min_df=10, max_df=0.1, token_pattern=r'[^ ]+')),\n",
+    "        (\"vectorizer\", TfidfVectorizer(min_df=10, max_df=0.1, token_pattern=r\"[^ ]+\")),\n",
     "        (\"classifier\", ComplementNB(alpha=0.5, fit_prior=False)),\n",
     "    ],\n",
     "    verbose=1,\n",
@@ -405,6 +407,7 @@
    "source": [
     "import logging\n",
     "import sys\n",
+    "\n",
     "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
     "\n",
     "save_model(classifier, key=model_key, keep_last_n=1)"

good_ai/src/good_ai/__init__.py

@@ -1,3 +1,3 @@
 from .good_ai import *
-from .utilities import *
 from .open_s3 import *
+from .utilities import *

good_ai/src/good_ai/good_ai/models/load_model.py

@@ -2,6 +2,7 @@ import logging
 from typing import Any, Optional
 
 from joblib import load
+
 from good_ai.open_s3 import LargeFile
 
 logger = logging.getLogger("models")

good_ai/src/good_ai/good_ai/models/save_model.py

@@ -3,6 +3,7 @@ from pathlib import Path
 from typing import Optional, Union
 
 from joblib import dump
+
 from good_ai.open_s3 import LargeFile
 
 logger = logging.getLogger("models")

good_ai/src/good_ai/utilities/publication_tei/publication_tei.py

@@ -3,6 +3,7 @@ from typing import Any, List, Optional, Union
 
 from bs4 import BeautifulSoup
 from bs4.element import NavigableString, Tag
+
 from good_ai.utilities.publication_tei.models.element import Paragraph
 
 from ..clean import clean