Update app.py

app.py

@@ -32,14 +32,23 @@ def init_zsl_topic_classification():
     pipe = pipeline("zero-shot-classification", model=MODEL)
     template = "This text is about {}."
     return pipe, template
+
+@st.cache(allow_output_mutation = True)
+def init_ner_pipeline():
+    tokenizer = AutoTokenizer.from_pretrained("d4data/biomedical-ner-all")
+    model = AutoModelForTokenClassification.from_pretrained("d4data/biomedical-ner-all")
+    pipe = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple") # pass device=0 if using gpu
+    return pipe
+
     
 # Model initialization    
 pipeline_summarization = init_text_summarization_model()
 pipeline_zsl, template = init_zsl_topic_classification()
+pipeline_ner =init_ner_pipeline()
 
 st.header("Intelligent Document Automation")
 
-def get_text_from_ocr_engine(uploaded_file):
+def get_text_from_ocr_engine():
     return "This is a sample text for named entity recognition and other tasks"
 
 with st.sidebar:
@@ -52,4 +61,5 @@ if selected_menu == "Upload Document":
     if uploaded_file is not None:
         ocr_text  = get_text_from_ocr_engine(uploaded_file)
         st.write("Upload Successful")
-     
+elif selected_menu == "Upload Document:
+    st.write(get_text_from_ocr_engine())