Update interfaces/base_pipeline.py

interfaces/base_pipeline.py

@@ -1,42 +1,42 @@
-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."""
-    })
-
+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(device="cpu", 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)