From 2ed13e715b41acaf6207bf74463c7098de24b6a0 Mon Sep 17 00:00:00 2001 From: Roman Kalyakin Date: Wed, 23 Oct 2024 11:22:04 +0200 Subject: [PATCH] make mypy happy --- impresso/resources/tools.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/impresso/resources/tools.py b/impresso/resources/tools.py index fb1c1c0..91e4bd7 100644 --- a/impresso/resources/tools.py +++ b/impresso/resources/tools.py @@ -3,6 +3,9 @@ from impresso.api_client.models.impresso_named_entity_recognition_request import ( ImpressoNamedEntityRecognitionRequest, ) +from impresso.api_client.models.impresso_named_entity_recognition_request_method import ( + ImpressoNamedEntityRecognitionRequestMethod, +) from impresso.api_models import ImpressoNerResponse from impresso.data_container import DataContainer from impresso.resources.base import Resource @@ -65,7 +68,9 @@ def ner(self, text: str) -> NerContainer: """ result = perform_ner.sync( client=self._api_client, - body=ImpressoNamedEntityRecognitionRequest(text=text, method="ner"), + body=ImpressoNamedEntityRecognitionRequest( + text=text, method=ImpressoNamedEntityRecognitionRequestMethod.NER + ), ) raise_for_error(result) @@ -88,7 +93,9 @@ def ner_nel(self, text: str) -> NerContainer: """ result = perform_ner.sync( client=self._api_client, - body=ImpressoNamedEntityRecognitionRequest(text=text, method="ner-nel"), + body=ImpressoNamedEntityRecognitionRequest( + text=text, method=ImpressoNamedEntityRecognitionRequestMethod.NER_NEL + ), ) raise_for_error(result) @@ -111,7 +118,9 @@ def nel(self, text: str) -> NerContainer: """ result = perform_ner.sync( client=self._api_client, - body=ImpressoNamedEntityRecognitionRequest(text=text, method="nel"), + body=ImpressoNamedEntityRecognitionRequest( + text=text, method=ImpressoNamedEntityRecognitionRequestMethod.NEL + ), ) raise_for_error(result)