Update interfaces/base_pipeline.py

interfaces/base_pipeline.py

@@ -2,7 +2,7 @@ from utca.core import RenameAttribute,Flush
 from utca.implementation.predictors import TokenSearcherPredictor, TokenSearcherPredictorConfig
 from utca.implementation.tasks import TokenSearcherNER, TokenSearcherNERPostprocessor
 from utca.implementation.predictors.token_searcher.token_searcher_pipeline import TokenClassificationPipeline
-from transformers import AutoTokenizer, AutoModelForTokenClassification
+from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
 
 predictor = TokenSearcherPredictor(
     TokenSearcherPredictorConfig(
@@ -29,7 +29,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  = TokenClassificationPipeline(device="cpu", model=model, tokenizer=tokenizer, aggregation_strategy = 'first')
+transformers_pipeline  = pipeline(task="ner", model ="knowledgator/UTC-DeBERTa-small", pipeline_class =TokenClassificationPipeline, aggregation_strategy = "first", batch_size = 12)
 
 if __name__=="__main__":
     pipeline = generate_pipeline()