import streamlit as st from Final_file import FlairRecognizer from Final_file import FlairRecognizer2 import os import PyPDF2 import docx # from io import BytesIO from fpdf import FPDF import io from docx import Document from PiiMaskingService import PiiMaskingService # Cache the model loading and prediction function @st.cache_resource def cached_predict_ner_tags(text): return FlairRecognizer.predict_ner_tags(text) # Cache the text analysis function @st.cache_resource def cached_analyze_text(text, operator): return FlairRecognizer.analyze_text(text) @st.cache_resource def cached_anonimize_text(text, operator): return FlairRecognizer2.anonymize(text, operator) @st.cache_resource def anonymize(text, operator, model): return PiiMaskingService().anonymize(text, operator, model) def download_masked_file(masked_text, file_extension): # Create a temporary file to store the masked text temp_file_path = f"masked_output.{file_extension}" with open(temp_file_path, "w") as temp_file: temp_file.write(masked_text) # Display a download button st.download_button("Download Masked File", temp_file_path, file_name=f"masked_output.{file_extension}") # Clean up the temporary file os.remove(temp_file_path) def extract_text_from_pdf(file_contents): try: # base64_pdf = base64.b64encode(file_contents.read()).decode('utf-8') pdf_reader = PyPDF2.PdfReader(file_contents) text = '' for page_num in range(len(pdf_reader.pages)): text += pdf_reader.pages[page_num].extract_text() return text except Exception as e: return f"Error occurred: {str(e)}" def create_pdf(text_content): pdf = FPDF() pdf.add_page() pdf.add_font("DejaVuSans", "", "DejaVuSans.ttf",uni=True) # Add DejaVuSans font pdf.set_font("DejaVuSans", size=12) pdf.multi_cell(0, 10, txt=text_content) return pdf def create_word_file(text_content): doc = Document() doc.add_paragraph(text_content) # Save the document to a BytesIO object doc_io = io.BytesIO() doc.save(doc_io) doc_io.seek(0) return doc_io def main(): st.title('PII Masking App') st.sidebar.header('Upload Options') upload_option = st.sidebar.radio("Choose upload option:", ('Text Input', 'File Upload')) st_operator = st.sidebar.selectbox( "De-identification approach", ["redact", "replace", "encrypt", "hash", "mask"], index=1, help=""" Select which manipulation to the text is requested after PII has been identified.\n - Redact: Completely remove the PII text\n - Replace: Replace the PII text with a constant, e.g. \n - Highlight: Shows the original text with PII highlighted in colors\n - Mask: Replaces a requested number of characters with an asterisk (or other mask character)\n - Hash: Replaces with the hash of the PII string\n - Encrypt: Replaces with an AES encryption of the PII string, allowing the process to be reversed """, ) st_model = st.sidebar.selectbox( "NER model package", [ "flair/ner-english-large", "HuggingFace/obi/deid_roberta_i2b2", ], index=0, ) masked_text_public = '' if upload_option == 'Text Input': input_text = st.text_area("Enter text here:") if st.button('Analyze'): with st.spinner('Wait for it... the model is loading'): # cached_predict_ner_tags(input_text) masked_text = anonymize(input_text, st_operator, st_model) # masked_text = cached_anonimize_text(input_text, st_operator) st.text_area("Masked text:", value=masked_text, height=200) elif upload_option == 'File Upload': uploaded_file = st.file_uploader("Upload a file", type=['txt', 'pdf', 'docx']) if uploaded_file is not None: file_contents = uploaded_file.read() # Process PDF file if uploaded_file.type == 'application/pdf': extracted_text = extract_text_from_pdf(uploaded_file) if st.button('Analyze'): with st.spinner('Wait for it... the model is loading'): # cached_predict_ner_tags(extracted_text) masked_text = anonymize(extracted_text, st_operator, st_model) # masked_text = cached_analyze_text(extracted_text) st.text_area("Masked text:", value=masked_text, height=200) # Display the extracted text if extracted_text: pdf = create_pdf(masked_text) # Save PDF to temporary location pdf_file_path = "masked_output.pdf" pdf.output(pdf_file_path) # Download button st.download_button(label="Download", data=open(pdf_file_path, "rb"), file_name="masked_output.pdf", mime="application/pdf") else: st.warning("Please enter some text to download as PDF.") # Process Word document elif uploaded_file.type == 'application/vnd.openxmlformats-officedocument.wordprocessingml.document': doc = docx.Document(io.BytesIO(file_contents)) text = '' for paragraph in doc.paragraphs: text += paragraph.text if st.button('Analyze'): with st.spinner('Wait for it... the model is loading'): # cached_predict_ner_tags(text) masked_text = anonymize(text, st_operator, st_model) # masked_text = cached_analyze_text(text) st.text_area("Masked text:", value=masked_text, height=200) #create word file doc_io = create_word_file(masked_text) #download it st.download_button(label="Download", data=doc_io, file_name="masked_text.docx", mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document") else: if st.button('Analyze'): with st.spinner('Wait for it... the model is loading'): # cached_predict_ner_tags(file_contents.decode()) # masked_text = cached_analyze_text(file_contents.decode()) masked_text = anonymize(file_contents.decode(), st_operator, st_model) st.text_area("Masked text:", value=masked_text, height=200) st.download_button(label="Download",data = masked_text,file_name="masked_text.txt") if __name__ == "__main__": main()