Update interfaces/base_pipeline.py

interfaces/base_pipeline.py

@@ -1,3 +1,4 @@
+import torch
 from utca.core import RenameAttribute,Flush
 from utca.implementation.predictors import TokenSearcherPredictor, TokenSearcherPredictorConfig
 from utca.implementation.tasks import TokenSearcherNER, TokenSearcherNERPostprocessor
@@ -29,7 +30,7 @@ def generate_pipeline(threshold: float = 0.5):
 tokenizer = AutoTokenizer.from_pretrained("knowledgator/UTC-DeBERTa-large-v2")
 model = AutoModelForTokenClassification.from_pretrained("knowledgator/UTC-DeBERTa-large-v2")
 
-transformers_pipeline = pipeline(device="cpu", task="ner", model ="knowledgator/UTC-DeBERTa-small", pipeline_class =TokenClassificationPipeline, aggregation_strategy = "first", batch_size = 12)
+transformers_pipeline = pipeline(device=torch.device("cpu"), task="ner", model ="knowledgator/UTC-DeBERTa-small", pipeline_class =TokenClassificationPipeline, aggregation_strategy = "first", batch_size = 12)
 
 if __name__=="__main__":
  pipeline = generate_pipeline()