Improve example

This commit is contained in:
Andras Schmelczer 2022-05-26 21:49:58 +02:00
parent 3e1564b9d0
commit b8a373fa80
No known key found for this signature in database
GPG key ID: 39260B5B0614A13E
2 changed files with 20 additions and 23 deletions

View file

@ -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]