From 8033e63b51a4687b96d2353099d4ce7343d0d53b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A1s=20Schmelczer?= <andras@schmelczer.dev>
Date: Thu, 26 May 2022 21:49:58 +0200
Subject: [PATCH] Improve example

---
 examples/simple/helpers.py        |  2 +-
 examples/simple/predict_domain.py | 41 ++++++++++++++-----------------
 2 files changed, 20 insertions(+), 23 deletions(-)

diff --git a/examples/simple/helpers.py b/examples/simple/helpers.py
index af2c51e..923fd19 100644
--- a/examples/simple/helpers.py
+++ b/examples/simple/helpers.py
@@ -12,5 +12,5 @@ def preprocess(text: str) -> str:
 
 def lemmatize(text: str) -> str:
     lemmatized = lemmatize_text(text)
-    clean_lemmas = [re.sub(r"\d[\d.,]*", "NUM", lemma) for lemma in lemmatized]
+    clean_lemmas = [re.sub(r"\d[\d.,]*", "NUM", lemma.lower()) for lemma in lemmatized]
     return " ".join(clean_lemmas)
diff --git a/examples/simple/predict_domain.py b/examples/simple/predict_domain.py
index 939114f..edfa662 100644
--- a/examples/simple/predict_domain.py
+++ b/examples/simple/predict_domain.py
@@ -1,5 +1,6 @@
-from typing import Dict, Iterable, List
-from great_ai import GreatAI, use_model, ClassificationOutput
+from typing import Iterable, List, Optional
+
+from great_ai import ClassificationOutput, GreatAI, use_model
 from sklearn.pipeline import Pipeline
 
 from helpers import lemmatize, preprocess
@@ -7,39 +8,37 @@ from helpers import lemmatize, preprocess
 
 @GreatAI.deploy
 @use_model("small-domain-prediction-v2", version="latest")
-def predict_domain(text: str, model: Pipeline, target_confidence: float = 20) -> List[ClassificationOutput]:
+def predict_domain(
+    text: str, model: Pipeline, target_confidence: float = 20
+) -> List[ClassificationOutput]:
     """
     Predict the scientific domain of the input text.
     Return labels until their sum likelihood is larger than target_confidence.
     """
     assert 0 <= target_confidence <= 100, "invalid argument"
 
-    text = preprocess(text)
+    processed = [(word, lemmatize(word)) for word in preprocess(text).split(" ")]
+    token_mapping = dict(processed)
+    clean_input = " ".join(v[1] for v in processed)
 
-    token_mapping = {lemmatize(original): original for original in text.split(" ")}
     feature_names = [
         token_mapping.get(name)
         for name in model.named_steps["vectorizer"].get_feature_names_out()
     ]
 
-    features = model.named_steps["vectorizer"].transform(
-        [" ".join(token_mapping.keys())]
-    )
-    prediction = model.named_steps["classifier"].predict_proba(features)[0]
+    prediction = model.predict_proba([clean_input])[0]
     best_classes = sorted(enumerate(prediction), key=lambda v: v[1], reverse=True)
 
     results: List[ClassificationOutput] = []
     for class_index, probability in best_classes:
-        weights = model.named_steps["classifier"].feature_log_prob_[class_index]
-        domain = model.named_steps["classifier"].classes_[class_index]
-
         results.append(
             ClassificationOutput(
-                label=domain,
+                label=model.named_steps["classifier"].classes_[class_index],
                 confidence=round(probability * 100),
                 explanation=_get_explanation(
-                    weights=weights,
-                    counts=features.A[0],
+                    weights=model.named_steps["classifier"].feature_log_prob_[
+                        class_index
+                    ],
                     words=feature_names,
                 ),
             )
@@ -53,13 +52,11 @@ def predict_domain(text: str, model: Pipeline, target_confidence: float = 20) ->
 
 def _get_explanation(
     weights: Iterable[float],
-    counts: Iterable[float],
-    words: Iterable[str],
+    words: Iterable[Optional[str]],
 ) -> List[str]:
-    most_influential = sorted((
-        (weight, word)
-        for weight, count, word in zip(weights, counts, words)
-        if count > 0
-    ), reverse=True)[:5]
+    most_influential = sorted(
+        ((weight, word) for weight, word in zip(weights, words) if word),
+        reverse=True,
+    )[:5]
 
     return [word for _, word in most_influential]