62 lines
1.9 KiB
Python
62 lines
1.9 KiB
Python
from typing import Iterable, List, Optional
|
|
|
|
from great_ai import ClassificationOutput, GreatAI, use_model
|
|
from sklearn.pipeline import Pipeline
|
|
|
|
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]:
|
|
"""
|
|
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"
|
|
|
|
processed = [(word, lemmatize(word)) for word in preprocess(text).split(" ")]
|
|
token_mapping = dict(processed)
|
|
clean_input = " ".join(v[1] for v in processed)
|
|
|
|
feature_names = [
|
|
token_mapping.get(name)
|
|
for name in model.named_steps["vectorizer"].get_feature_names_out()
|
|
]
|
|
|
|
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:
|
|
results.append(
|
|
ClassificationOutput(
|
|
label=model.named_steps["classifier"].classes_[class_index],
|
|
confidence=round(probability * 100),
|
|
explanation=_get_explanation(
|
|
weights=model.named_steps["classifier"].feature_log_prob_[
|
|
class_index
|
|
],
|
|
words=feature_names,
|
|
),
|
|
)
|
|
)
|
|
|
|
if sum(r.confidence for r in results) >= target_confidence:
|
|
break
|
|
|
|
return results
|
|
|
|
|
|
def _get_explanation(
|
|
weights: Iterable[float],
|
|
words: Iterable[Optional[str]],
|
|
) -> List[str]:
|
|
most_influential = sorted(
|
|
((weight, word) for weight, word in zip(weights, words) if word),
|
|
reverse=True,
|
|
)[:5]
|
|
|
|
return [word for _, word in most_influential]
|