Format
Signed-off-by: András Schmelczer <andras@schmelczer.dev>
This commit is contained in:
parent
acbdd96883
commit
1c501db81a
8 changed files with 16 additions and 11 deletions
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue