Update app.py

app.py

@@ -1,7 +1,8 @@
-from transformers import AutoTokenizer
-from transformers import AutoModelForSeq2SeqLM
+from transformers import pipeline
 import streamlit as st
 import fitz  # PyMuPDF
+from transformers import AutoTokenizer
+from transformers import AutoModelForSeq2SeqLM
 from docx import Document
 import re
 import nltk
@@ -11,24 +12,11 @@ def sentence_tokenize(text):
     sentences = nltk.sent_tokenize(text)
     return sentences
 
+# Use a pipeline as a high-level helper
 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 + '.'
@@ -42,54 +30,7 @@ def mask_generation(text, model=model_large, tokenizer=tokenizer_large):
     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
+pipe1 = pipeline("token-classification", model="edithram23/new-bert-v2")
 
 def redact_text(page, text):
     text_instances = page.search_for(text)
@@ -132,37 +73,17 @@ if uploaded_file is not None:
     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)
+              final=[]
+              text = pg.get_text()
+              sentences = sentence_tokenize(text)
+              for sentence in sentences:
+                x=[pipe1(sentence)]
+                m = combine_words(x[0])
+                for j in m:
+                  if(j['entity']!='none'):
+                    final.append(j['word'])
+              for i in final:
+                redact_text(pg,i)
         output_pdf = "output_redacted.pdf"
         pdf_document.save(output_pdf)