app.py

from transformers import AutoTokenizer
from transformers import AutoModelForSeq2SeqLM
import streamlit as st
import fitz  # PyMuPDF
from docx import Document
import re
import nltk
nltk.download('punkt')

def sentence_tokenize(text):
    sentences = nltk.sent_tokenize(text)
    return sentences

model_dir_large = 'edithram23/Redaction_Personal_info_v1'
tokenizer_large = AutoTokenizer.from_pretrained(model_dir_large)
model_large = AutoModelForSeq2SeqLM.from_pretrained(model_dir_large)

# model_dir_small = 'edithram23/Redaction'
# tokenizer_small = AutoTokenizer.from_pretrained(model_dir_small)
# model_small = AutoModelForSeq2SeqLM.from_pretrained(model_dir_small)

# def small(text, model=model_small, tokenizer=tokenizer_small):
#     inputs = ["Mask Generation: " + text.lower() + '.']
#     inputs = tokenizer(inputs, max_length=256, truncation=True, return_tensors="pt")
#     output = model.generate(**inputs, num_beams=8, do_sample=True, max_length=len(text))
#     decoded_output = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
#     predicted_title = decoded_output.strip()
#     pattern = r'\[.*?\]'
#     redacted_text = re.sub(pattern, '[redacted]', predicted_title)
#     return redacted_text

def mask_generation(text, model=model_large, tokenizer=tokenizer_large):
    if len(text) < 90:
        text = text + '.'
        # return small(text)
    inputs = ["Mask Generation: " + text.lower() + '.']
    inputs = tokenizer(inputs, max_length=512, truncation=True, return_tensors="pt")
    output = model.generate(**inputs, num_beams=8, do_sample=True, max_length=len(text))
    decoded_output = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
    predicted_title = decoded_output.strip()
    pattern = r'\[.*?\]'
    redacted_text = re.sub(pattern, '[redacted]', predicted_title)
    return redacted_text

def find_surrounding_words(text, target="[redacted]"):
    pattern = re.compile(r'([A-Za-z0-9_@#\$%\^&*\(\)\[\]\{\}\.\,]+)?\s*' + re.escape(target) + r'\s*([A-Za-z0-9_@#\$%\^&*\(\)\[\]\{\}\.\,]+)?')
    matches = pattern.finditer(text)
    results = []
    for match in matches:
        before, after = match.group(1), match.group(2)

        if before:
            before_parts = before.split(',')
            before_parts = [item for item in before_parts if item.strip()]
            if len(before_parts) > 1:
                before_word = before_parts[0].strip()
                before_index = match.start(1)
            else:
                before_word = before_parts[0]
                before_index = match.start(1)
        else:
            before_word = None
            before_index = None

        if after:
            after_parts = after.split(',')
            after_parts = [item for item in after_parts if item.strip()]
            if len(after_parts) > 1:
                after_word = after_parts[0].strip()
                after_index = match.start(2)
            else:
                after_word = after_parts[0]
                after_index = match.start(2)
        else:
            after_word = None
            after_index = None

        if match.start() == 0:
            before_word = None
            before_index = None

        if match.end() == len(text):
            after_word = None
            after_index = None

        results.append({
            "before_word": before_word,
            "after_word": after_word,
            "before_index": before_index,
            "after_index": after_index
        })
    return results

def redact_text(page, text):
    text_instances = page.search_for(text)
    for inst in text_instances:
        page.add_redact_annot(inst, fill=(0, 0, 0))
    page.apply_redactions()

def read_pdf(file):
    pdf_document = fitz.open(stream=file.read(), filetype="pdf")
    text = ""
    for page_num in range(len(pdf_document)):
        page = pdf_document.load_page(page_num)
        text += page.get_text()
    return text, pdf_document

def read_docx(file):
    doc = Document(file)
    text = "\n".join([para.text for para in doc.paragraphs])
    return text

def read_txt(file):
    text = file.read().decode("utf-8")
    return text

def process_file(file):
    if file.type == "application/pdf":
        return read_pdf(file)
    elif file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
        return read_docx(file), None
    elif file.type == "text/plain":
        return read_txt(file), None
    else:
        return "Unsupported file type.", None

st.title("Redaction")
uploaded_file = st.file_uploader("Upload a file", type=["pdf", "docx", "txt"])

if uploaded_file is not None:
    file_contents, pdf_document = process_file(uploaded_file)
    if pdf_document:
        redacted_text = []
        for page in pdf_document:
            pg = page.get_text()
            pg_lower = pg.lower()
            token = sentence_tokenize(pg)
            final = ''
            for t in token:
                t_lower = t.lower()
                final = mask_generation(t)
                words = find_surrounding_words(final)
                for i in range(len(words)):
                    if words[i]['after_index'] is None:
                        if words[i]['before_word'] in t_lower:
                            fi = t_lower.index(words[i]['before_word'])
                            fi = fi + len(words[i]['before_word'])
                            li = len(t)
                            redacted_text.append(t[fi:li])
                    elif words[i]['before_index'] is None:
                        if words[i]['after_word'] in t_lower:
                            fi = 0
                            li = t_lower.index(words[i]['after_word'])
                            redacted_text.append(t[fi:li])
                    else:
                        if words[i]['after_word'] in t_lower and words[i]['before_word'] in t_lower:
                            before_word = words[i]['before_word']
                            after_word = words[i]['after_word']
                            fi = t_lower.index(before_word)
                            fi = fi + len(before_word)
                            li = t_lower.index(after_word)
                            redacted_text.append(t[fi:li])
        for page in pdf_document:
            for i in redacted_text:
                redact_text(page, i)
        output_pdf = "output_redacted.pdf"
        pdf_document.save(output_pdf)

        with open(output_pdf, "rb") as file:
            st.download_button(
                label="Download Processed PDF",
                data=file,
                file_name="processed_file.pdf",
                mime="application/pdf",
            )
    else:
        token = sentence_tokenize(file_contents)
        final = ''
        for i in range(0, len(token)):
            final += mask_generation(token[i]) + '\n'
        processed_text = final
        st.text_area("OUTPUT", processed_text, height=400)
        st.download_button(
            label="Download Processed File",
            data=processed_text,
            file_name="processed_file.txt",
            mime="text/plain",
        )