OnlyBiggg commited on
Commit
3987b72
·
1 Parent(s): 5816e57

fix accelerator

Browse files
Files changed (1) hide show
  1. app/ner/services/ner.py +4 -1
app/ner/services/ner.py CHANGED
@@ -14,7 +14,10 @@ class NER:
14
 
15
  self.tokenizer = AutoTokenizer.from_pretrained(self.model_dir)
16
  self.model = ORTModelForTokenClassification.from_pretrained(self.model_dir)
17
- self.pipeline = pipeline(task=settings.TASK_NAME, model=self.model, tokenizer=self.tokenizer, accelerator=settings.ACCELERATOR, device=settings.DEVICE)
 
 
 
18
 
19
  def predict(self, text: str, entity_tag: str = None):
20
  if self.pipeline is None:
 
14
 
15
  self.tokenizer = AutoTokenizer.from_pretrained(self.model_dir)
16
  self.model = ORTModelForTokenClassification.from_pretrained(self.model_dir)
17
+ self.pipeline = pipeline(task=settings.TASK_NAME,
18
+ model=self.model,
19
+ tokenizer=self.tokenizer,
20
+ device=settings.DEVICE)
21
 
22
  def predict(self, text: str, entity_tag: str = None):
23
  if self.pipeline is None: