Improve example
This commit is contained in:
parent
3e1564b9d0
commit
b8a373fa80
2 changed files with 20 additions and 23 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue