Improve example

This commit is contained in:
Andras Schmelczer 2022-05-26 21:49:58 +02:00
parent e4b4c262c3
commit 8033e63b51
2 changed files with 20 additions and 23 deletions

View file

@ -12,5 +12,5 @@ def preprocess(text: str) -> str:
def lemmatize(text: str) -> str: def lemmatize(text: str) -> str:
lemmatized = lemmatize_text(text) 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) return " ".join(clean_lemmas)

View file

@ -1,5 +1,6 @@
from typing import Dict, Iterable, List from typing import Iterable, List, Optional
from great_ai import GreatAI, use_model, ClassificationOutput
from great_ai import ClassificationOutput, GreatAI, use_model
from sklearn.pipeline import Pipeline from sklearn.pipeline import Pipeline
from helpers import lemmatize, preprocess from helpers import lemmatize, preprocess
@ -7,39 +8,37 @@ from helpers import lemmatize, preprocess
@GreatAI.deploy @GreatAI.deploy
@use_model("small-domain-prediction-v2", version="latest") @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. Predict the scientific domain of the input text.
Return labels until their sum likelihood is larger than target_confidence. Return labels until their sum likelihood is larger than target_confidence.
""" """
assert 0 <= target_confidence <= 100, "invalid argument" 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 = [ feature_names = [
token_mapping.get(name) token_mapping.get(name)
for name in model.named_steps["vectorizer"].get_feature_names_out() for name in model.named_steps["vectorizer"].get_feature_names_out()
] ]
features = model.named_steps["vectorizer"].transform( prediction = model.predict_proba([clean_input])[0]
[" ".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) best_classes = sorted(enumerate(prediction), key=lambda v: v[1], reverse=True)
results: List[ClassificationOutput] = [] results: List[ClassificationOutput] = []
for class_index, probability in best_classes: 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( results.append(
ClassificationOutput( ClassificationOutput(
label=domain, label=model.named_steps["classifier"].classes_[class_index],
confidence=round(probability * 100), confidence=round(probability * 100),
explanation=_get_explanation( explanation=_get_explanation(
weights=weights, weights=model.named_steps["classifier"].feature_log_prob_[
counts=features.A[0], class_index
],
words=feature_names, words=feature_names,
), ),
) )
@ -53,13 +52,11 @@ def predict_domain(text: str, model: Pipeline, target_confidence: float = 20) ->
def _get_explanation( def _get_explanation(
weights: Iterable[float], weights: Iterable[float],
counts: Iterable[float], words: Iterable[Optional[str]],
words: Iterable[str],
) -> List[str]: ) -> List[str]:
most_influential = sorted(( most_influential = sorted(
(weight, word) ((weight, word) for weight, word in zip(weights, words) if word),
for weight, count, word in zip(weights, counts, words) reverse=True,
if count > 0 )[:5]
), reverse=True)[:5]
return [word for _, word in most_influential] return [word for _, word in most_influential]