Signed-off-by: András Schmelczer <andras@schmelczer.dev>
This commit is contained in:
Andras Schmelczer 2022-04-02 21:59:09 +02:00
parent acbdd96883
commit 1c501db81a
8 changed files with 16 additions and 11 deletions

View file

@ -2,13 +2,14 @@ import re
from typing import Dict, Iterable, List
from config import model_key
from preprocess import preprocess
from models import DomainPrediction
from preprocess import preprocess
from sklearn.pipeline import Pipeline
from good_ai import use_model
from good_ai.utilities.clean import clean
@use_model(model_key, version="latest")
def predict_domain(
text: str, model: Pipeline, cut_off_probability: float = 0.2
@ -19,12 +20,11 @@ def predict_domain(
feature_names = model.named_steps["vectorizer"].get_feature_names_out()
token_mapping = {
preprocess(original): original
for original in text.split(" ")
}
token_mapping = {preprocess(original): original for original in text.split(" ")}
features = model.named_steps["vectorizer"].transform([" ".join(token_mapping.keys())])
features = model.named_steps["vectorizer"].transform(
[" ".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)