interfaces/base_pipeline.py

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 
 
predictor = TokenSearcherPredictor( 
 TokenSearcherPredictorConfig( 
 model_name = "knowledgator/UTC-DeBERTa-large-v2", 
 device="cpu" 
 ) 
) 
 
def generate_pipeline(threshold: float = 0.5): 
 ner_task = TokenSearcherNER( 
 predictor=predictor, 
 postprocess=TokenSearcherNERPostprocessor( 
 threshold=threshold 
 ) 
 ) 
 pipeline = ( 
 ner_task 
 | Flush(keys=["labels"]) 
 | RenameAttribute("output", "entities") 
 ) 
 
 return pipeline 
 
tokenizer = AutoTokenizer.from_pretrained("knowledgator/UTC-DeBERTa-large-v2") 
model = AutoModelForTokenClassification.from_pretrained("knowledgator/UTC-DeBERTa-large-v2") 
 
transformers_pipeline = TokenClassificationPipeline(model=model, tokenizer=tokenizer, aggregation_strategy = 'first') 
 
if __name__=="__main__": 
 pipeline = generate_pipeline() 
 res = pipeline.run({ 
 "text": """Dr. Paul Hammond, a renowned neurologist at Johns Hopkins University, has recently published a paper in the prestigious journal "Nature Neuroscience". 
 His research focuses on a rare genetic mutation, found in less than 0.01% of the population, that appears to prevent the development of Alzheimer's disease. Collaborating with researchers at the University of California, San Francisco, the team is now working to understand the mechanism by which this mutation confers its protective effect. 
 Funded by the National Institutes of Health, their research could potentially open new avenues for Alzheimer's treatment.""" 
 }) 
 
 print(res)