edithram23/Redaction_PDF
0
1from transformers import pipeline2from transformers import AutoTokenizer3from transformers import AutoModelForSeq2SeqLM4import streamlit as st5import fitz # PyMuPDF6from docx import Document7import re8import nltk9from presidio_analyzer import AnalyzerEngine, PatternRecognizer, RecognizerResult, Pattern10nltk.download('punkt')11 12 13def sentence_tokenize(text):14 sentences = nltk.sent_tokenize(text)15 return sentences16 17model_dir_large = 'edithram23/Redaction_Personal_info_v1'18tokenizer_large = AutoTokenizer.from_pretrained(model_dir_large)19model_large = AutoModelForSeq2SeqLM.from_pretrained(model_dir_large)20pipe1 = pipeline("token-classification", model="edithram23/new-bert-v2")21 22# model_dir_small = 'edithram23/Redaction'23# tokenizer_small = AutoTokenizer.from_pretrained(model_dir_small)24# model_small = AutoModelForSeq2SeqLM.from_pretrained(model_dir_small)25 26# def small(text, model=model_small, tokenizer=tokenizer_small):27# inputs = ["Mask Generation: " + text.lower() + '.']28# inputs = tokenizer(inputs, max_length=256, truncation=True, return_tensors="pt")29# output = model.generate(**inputs, num_beams=8, do_sample=True, max_length=len(text))30# decoded_output = tokenizer.batch_decode(output, skip_special_tokens=True)[0]31# predicted_title = decoded_output.strip()32# pattern = r'\[.*?\]'33# redacted_text = re.sub(pattern, '[redacted]', predicted_title)34# return redacted_text35 36# Initialize the analyzer engine37analyzer = AnalyzerEngine()38 39# Define a custom address recognizer using a regex pattern40address_pattern = Pattern(name="address", regex=r"\d+\s\w+\s(?:street|st|road|rd|avenue|ave|lane|ln|drive|dr|blvd|boulevard)\s*\w*", score=0.5)41address_recognizer = PatternRecognizer(supported_entity="ADDRESS", patterns=[address_pattern])42 43# Add the custom address recognizer to the analyzer44analyzer.registry.add_recognizer(address_recognizer)45# analyzer.get_recognizers46# Define a function to extract entities47 48 49def combine_words(entities):50 combined_entities = []51 current_entity = None52 53 for entity in entities:54 if current_entity:55 if current_entity['end'] == entity['start']:56 # Combine the words without space57 current_entity['word'] += entity['word'].replace('##', '')58 current_entity['end'] = entity['end']59 elif current_entity['end'] + 1 == entity['start']:60 # Combine the words with a space61 current_entity['word'] += ' ' + entity['word'].replace('##', '')62 current_entity['end'] = entity['end']63 else:64 # Add the previous combined entity to the list65 combined_entities.append(current_entity)66 # Start a new entity67 current_entity = entity.copy()68 current_entity['word'] = current_entity['word'].replace('##', '')69 else:70 # Initialize the first entity71 current_entity = entity.copy()72 current_entity['word'] = current_entity['word'].replace('##', '')73 74 # Add the last entity75 if current_entity:76 combined_entities.append(current_entity)77 78 return combined_entities79 80def words_red_bert(text):81 final=[]82 sentences = sentence_tokenize(text)83 for sentence in sentences:84 x=[pipe1(sentence)]85 m = combine_words(x[0])86 for j in m:87 if(j['entity']!='none' and len(j['word'])>1 and j['word']!=', '):88 final.append(j['word'])89 return final90 91def extract_entities(text):92 entities = {93 "NAME": [],94 "PHONE_NUMBER": [],95 "EMAIL": [],96 "ADDRESS": [],97 "LOCATION": [],98 "IN_AADHAAR": [],99 }100 output = []101 102 # Analyze the text for PII103 results = analyzer.analyze(text=text, language='en')104 105 for result in results:106 if result.entity_type == "PERSON":107 entities["NAME"].append(text[result.start:result.end])108 output+=[text[result.start:result.end]]109 elif result.entity_type == "PHONE_NUMBER":110 entities["PHONE_NUMBER"].append(text[result.start:result.end])111 output+=[text[result.start:result.end]]112 elif result.entity_type == "EMAIL_ADDRESS":113 entities["EMAIL"].append(text[result.start:result.end])114 output+=[text[result.start:result.end]]115 elif result.entity_type == "ADDRESS":116 entities["ADDRESS"].append(text[result.start:result.end])117 output+=[text[result.start:result.end]]118 elif result.entity_type == 'LOCATION':119 entities['LOCATION'].append(text[result.start:result.end])120 output+=[text[result.start:result.end]]121 elif result.entity_type == 'IN_AADHAAR':122 entities['IN_PAN'].append(text[result.start:result.end])123 output+=[text[result.start:result.end]]124 125 return entities,output126 127def mask_generation(text, model=model_large, tokenizer=tokenizer_large):128 if len(text) < 90:129 text = text + '.'130 # return small(text)131 inputs = ["Mask Generation: " + text.lower() + '.']132 inputs = tokenizer(inputs, max_length=512, truncation=True, return_tensors="pt")133 output = model.generate(**inputs, num_beams=8, do_sample=True, max_length=len(text))134 decoded_output = tokenizer.batch_decode(output, skip_special_tokens=True)[0]135 predicted_title = decoded_output.strip()136 pattern = r'\[.*?\]'137 redacted_text = re.sub(pattern, '[redacted]', predicted_title)138 return redacted_text139 140def redact_text(page, text):141 text_instances = page.search_for(text)142 for inst in text_instances:143 page.add_redact_annot(inst, fill=(0, 0, 0))144 page.apply_redactions()145 146def read_pdf(file):147 pdf_document = fitz.open(stream=file.read(), filetype="pdf")148 text = ""149 for page_num in range(len(pdf_document)):150 page = pdf_document.load_page(page_num)151 text += page.get_text()152 return text, pdf_document153 154def read_docx(file):155 doc = Document(file)156 text = "\n".join([para.text for para in doc.paragraphs])157 return text158 159def read_txt(file):160 text = file.read().decode("utf-8")161 return text162 163def process_file(file):164 if file.type == "application/pdf":165 return read_pdf(file)166 elif file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":167 return read_docx(file), None168 elif file.type == "text/plain":169 return read_txt(file), None170 else:171 return "Unsupported file type.", None172 173st.title("Redaction")174uploaded_file = st.file_uploader("Upload a file", type=["pdf", "docx", "txt"])175 176if uploaded_file is not None:177 file_contents, pdf_document = process_file(uploaded_file)178 if pdf_document:179 redacted_text = ''180 for pg in pdf_document:181 text = pg.get_text()182 sentences = sentence_tokenize(text)183 for sent in sentences:184 # x = mask_generation(sent)185 186 # sent_n_q_c=[] 187 # sent_n = list(set(sent.lower().replace('.',' ').split("\n")))188 # for i in sent_n:189 # for j in i.split(" "):190 # sent_n_q_c+=j.split(',')191 # x_q = x.lower().replace('.',' ').split(' ') 192 # e=[]193 # for i in x_q:194 # e+=i.split(',') 195 # t5_words=set(sent_n_q_c).difference(set(e)) 196 entities,words_out = extract_entities(sent)197 # print("\nwords_out:",words_out)198 # print("\nT5",t5_words)199 # print("X:",x,"\nsent:",sent,"\nx_q:",x_q,"\nsent_n:",sent_n,"\ne:",e,"\nsent_n_q_c:",sent_n_q_c,'\nt5_words',t5_words)200 bert_words = words_red_bert(sent)201 # print("\nbert:",bert_words)202 new=[]203 for w in words_out:204 new+=w.split('\n')205 # words_out+=t5_words206 new+=bert_words207 words_out = [i for i in new if len(i)>3]208 # print("\nfinal:",words_out)209 words_out=sorted(words_out, key=len,reverse=True)210 211 for i in words_out:212 redact_text(pg,i)213 # st.text_area(redacted_text)214 215 output_pdf = "output_redacted.pdf"216 pdf_document.save(output_pdf)217 218 with open(output_pdf, "rb") as file:219 st.download_button(220 label="Download Processed PDF",221 data=file,222 file_name="processed_file.pdf",223 mime="application/pdf",224 )225 else:226 token = sentence_tokenize(file_contents)227 final = ''228 for i in range(0, len(token)):229 final += mask_generation(token[i]) + '\n'230 processed_text = final231 st.text_area("OUTPUT", processed_text, height=400)232 st.download_button(233 label="Download Processed File",234 data=processed_text,235 file_name="processed_file.txt",236 mime="text/plain",237 )238