Spaces:

srijaydeshpande
/

DeID

Running on Zero

App Files Files Community

srijaydeshpande commited on May 24, 2024

Commit

40e5815

verified ·

1 Parent(s): f41ea40

Update app.py

Browse files

Files changed (1) hide show

app.py +72 -56

app.py CHANGED Viewed

@@ -13,48 +13,49 @@ subprocess.run('pip install llama-cpp-agent==0.2.10', shell=True)
 from llama_cpp import Llama
-# HF_TOKEN = os.environ.get("HF_TOKEN", None)
 def process_document(pdf_path, page_ids=None):
-   extracted_pages = extract_pages(pdf_path, page_numbers=page_ids)
-   page2content = {}
-   for extracted_page in tqdm(extracted_pages):
-       page_id = extracted_page.pageid
-       content = process_page(extracted_page)
-       page2content[page_id] = content
-   return page2content
 def process_page(extracted_page):
-   content = []
-   elements = [element for element in extracted_page._objs]
-   elements.sort(key=lambda a: a.y1, reverse=True)
-   for i, element in enumerate(elements):
-       if isinstance(element, LTTextContainer):
-           line_text = extract_text_and_normalize(element)
-           content.append(line_text)
-   content = re.sub('\n+', ' ', ''.join(content))
-   return content
 def extract_text_and_normalize(element):
-   # Extract text from line and split it with new lines
-   line_texts = element.get_text().split('\n')
-   norm_text = ''
-   for line_text in line_texts:
-       line_text=line_text.strip()
-       if not line_text:
-           line_text = '\n'
-       else:
-           line_text = re.sub('\s+', ' ', line_text)
-           if not re.search('[\w\d\,\-]', line_text[-1]):
-               line_text+='\n'
-           else:
-               line_text+=' '
-       norm_text+=line_text
-   return norm_text
 def txt_to_html(text):
     html_content = "<html><body>"
@@ -63,23 +64,37 @@ def txt_to_html(text):
     html_content += "</body></html>"
     return html_content
-def deidentify_doc(pdftext="", prompt="", maxtokens=600, temperature=1.2, top_probability=0.95):
-    # prompt = "Please anonymize the following clinical note. Replace all the following information with the term '[redacted]': Redact any strings that might be a name or initials, patients’ names, doctors’ names, the names Dr., redact any medical staff names, redact any strings that might be a location or address, such as '3970 Longview Drive', redact any strings that look like 'age 37', redact any dates and registration numbers, redact professions such as 'manager', redact any contact information."
-    output = model.create_chat_completion(
-                    messages = [
-                        {"role": "assistant", "content": prompt},
-                        {
-                            "role": "user",
-                            "content": pdftext
-                        }
-                    ],
-                    max_tokens=maxtokens,
-                    temperature=temperature
-                )
-    output = output['choices'][0]['message']['content']
     return output
 def mkdir(dir):
@@ -88,19 +103,20 @@ def mkdir(dir):
 @spaces.GPU(duration=120)
 def pdf_to_text(files, output_folder, prompt, maxtokens=600, temperature=1.2, top_probability=0.95):
-    output_folder = output_folder.replace('\\','/')
     for file in files:
         file_name = os.path.basename(file)
         file_name_splt = file_name.split('.')
         print('File name is ', file_name)
         print('output folder is ', output_folder)
-        if(len(file_name_splt)>1 and file_name_splt[1]=='pdf'):
-           page2content = process_document(file, page_ids=[0])
-           pdftext = page2content[1]
-           if(pdftext):
-               anonymized_text = deidentify_doc(pdftext, prompt, maxtokens, temperature, top_probability)
     return anonymized_text
 model_id = "Meta-Llama-3-8B-Instruct.Q5_K_M.gguf"
 model = Llama(model_path=model_id, n_ctx=2048, n_threads=8, n_gpu_layers=81, n_batch=64)
@@ -113,11 +129,11 @@ input_folder_text = gr.Textbox(label='Enter output folder path')
 output_text = gr.Textbox()
 output_path_component = gr.File(label="Select Output Path")
 iface = gr.Interface(
-    fn = pdf_to_text,
-    inputs = ['files', input_folder_text, "textbox", max_tokens, temp_slider, prob_slider],
     outputs=output_text,
     title='COBIx Endoscopy Report De-Identification',
     description="This application assists to remove personal information from the uploaded clinical report",
     theme=gr.themes.Soft(),
-    )
 iface.launch()

 from llama_cpp import Llama
 def process_document(pdf_path, page_ids=None):
+    extracted_pages = extract_pages(pdf_path, page_numbers=page_ids)
+    page2content = {}
+    for extracted_page in tqdm(extracted_pages):
+        page_id = extracted_page.pageid
+        content = process_page(extracted_page)
+        page2content[page_id] = content
+    return page2content
 def process_page(extracted_page):
+    content = []
+    elements = [element for element in extracted_page._objs]
+    elements.sort(key=lambda a: a.y1, reverse=True)
+    for i, element in enumerate(elements):
+        if isinstance(element, LTTextContainer):
+            line_text = extract_text_and_normalize(element)
+            content.append(line_text)
+    content = re.sub('\n+', ' ', ''.join(content))
+    return content
 def extract_text_and_normalize(element):
+    # Extract text from line and split it with new lines
+    line_texts = element.get_text().split('\n')
+    norm_text = ''
+    for line_text in line_texts:
+        line_text = line_text.strip()
+        if not line_text:
+            line_text = '\n'
+        else:
+            line_text = re.sub('\s+', ' ', line_text)
+            if not re.search('[\w\d\,\-]', line_text[-1]):
+                line_text += '\n'
+            else:
+                line_text += ' '
+        norm_text += line_text
+    return norm_text
 def txt_to_html(text):
     html_content = "<html><body>"
     html_content += "</body></html>"
     return html_content
+def deidentify_doc(pdftext="", prompt="", maxtokens=600, temperature=1.2, top_probability=0.95):
+    def replace_words_with_asterisk(big_string, words_to_replace):
+        for word in words_to_replace:
+            big_string = big_string.replace(word, '*')
+        return big_string
+    def get_output(pdfcontent):
+        output = model.create_chat_completion(
+            messages=[
+                {"role": "assistant", "content": prompt},
+                {
+                    "role": "user",
+                    "content": pdfcontent
+                }
+            ],
+            max_tokens=maxtokens,
+            temperature=temperature
+        )
+        wordstoremove = output['choices'][0]['message']['content']
+        position = wordstoremove.find("STARTTOKEN,")
+        if position != -1:
+            wordstoremove = wordstoremove[position + len("STARTTOKEN,"):].strip()
+        output = replace_words_with_asterisk(pdftext, wordstoremove.split(','))
+        return output
+    iterations=2
+    output = pdftext
+    for _ in iterations:
+        output = get_output(output)
     return output
 def mkdir(dir):
 @spaces.GPU(duration=120)
 def pdf_to_text(files, output_folder, prompt, maxtokens=600, temperature=1.2, top_probability=0.95):
+    output_folder = output_folder.replace('\\', '/')
     for file in files:
         file_name = os.path.basename(file)
         file_name_splt = file_name.split('.')
         print('File name is ', file_name)
         print('output folder is ', output_folder)
+        if (len(file_name_splt) > 1 and file_name_splt[1] == 'pdf'):
+            page2content = process_document(file, page_ids=[0])
+            pdftext = page2content[1]
+            if (pdftext):
+                anonymized_text = deidentify_doc(pdftext, prompt, maxtokens, temperature, top_probability)
     return anonymized_text
 model_id = "Meta-Llama-3-8B-Instruct.Q5_K_M.gguf"
 model = Llama(model_path=model_id, n_ctx=2048, n_threads=8, n_gpu_layers=81, n_batch=64)
 output_text = gr.Textbox()
 output_path_component = gr.File(label="Select Output Path")
 iface = gr.Interface(
+    fn=pdf_to_text,
+    inputs=['files', input_folder_text, "textbox", max_tokens, temp_slider, prob_slider],
     outputs=output_text,
     title='COBIx Endoscopy Report De-Identification',
     description="This application assists to remove personal information from the uploaded clinical report",
     theme=gr.themes.Soft(),
+)
 iface.launch()