Skip to content
Snippets Groups Projects
Commit 55b31ce6 authored by Marcin Wątroba's avatar Marcin Wątroba
Browse files

Add metric calculator

parent 2f65ed87
No related merge requests found
......@@ -17,7 +17,7 @@ class WikineuralMultilingualNerTransformersWerProcessorBase(SentenceWerProcessor
self._tokenizer = AutoTokenizer.from_pretrained("Babelscape/wikineural-multilingual-ner")
self._model = AutoModelForTokenClassification.from_pretrained("Babelscape/wikineural-multilingual-ner").to(
'cuda')
self._pipeline = pipeline("ner", model=self._model, tokenizer=self._tokenizer, device='cuda')
self._pipeline = pipeline("ner", model=self._model, tokenizer=self._tokenizer, device='cuda:0')
def process_text(self, text: str) -> List[str]:
ner_results = self._pipeline(text)
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment