INNER CODE UNIT · Python
_predict
dragondrop-cloud/cloud-concierge · nlpengine/main.py:301
def _predict(
nlp: spacy.Language, new_resource_docs: dict, textcat_multilabel_model: Union
) -> dict:
"""
Predict the workspace category for each resource within `new_resource_docs`
using the trained `textcat_multilabel_model`.
"""
resource_name_to_workspace = {}
prediction_labels = textcat_multilabel_model.labels
for resource_name, doc in new_resource_docs.items():
current_document = nlp(doc)
scores_array = textcat_multilabel_model.predict([current_document])
resource_name_to_workspace[resource_name] = prediction_labels[
scores_array.argmax()
]
return resource_name_to_workspace