Update app.py

app.py

@@ -5,93 +5,9 @@ from transformers import pipeline, AutoModelForTokenClassification, AutoTokenize
 import PyPDF2
 import docx
 import io
+import re
 
-def chunk_text(text, chunk_size=128):
-    words = text.split()
-    chunks = []
-    current_chunk = []
-    current_length = 0
-
-    for word in words:
-        if current_length + len(word) + 1 > chunk_size:
-            chunks.append(' '.join(current_chunk))
-            current_chunk = [word]
-            current_length = len(word)
-        else:
-            current_chunk.append(word)
-            current_length += len(word) + 1
-
-    if current_chunk:
-        chunks.append(' '.join(current_chunk))
-
-    return chunks
-
-st.set_page_config(layout="wide")
-
-# Function to read text from uploaded file
-def read_file(file):
-    if file.type == "text/plain":
-        return file.getvalue().decode("utf-8")
-    elif file.type == "application/pdf":
-        pdf_reader = PyPDF2.PdfReader(io.BytesIO(file.getvalue()))
-        return " ".join(page.extract_text() for page in pdf_reader.pages)
-    elif file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
-        doc = docx.Document(io.BytesIO(file.getvalue()))
-        return " ".join(paragraph.text for paragraph in doc.paragraphs)
-    else:
-        st.error("Unsupported file type")
-        return None
-
-st.title("Turkish NER Models Testing")
-
-model_list = [
-    'girayyagmur/bert-base-turkish-ner-cased',
-    'savasy/bert-base-turkish-ner-cased',
-    'xlm-roberta-large-finetuned-conll03-english',
-    'asahi417/tner-xlm-roberta-base-ontonotes5'
-]
-
-st.sidebar.header("Select NER Model")
-model_checkpoint = st.sidebar.radio("", model_list)
-
-st.sidebar.write("For details of models: 'https://huggingface.co/akdeniz27/")
-st.sidebar.write("Only PDF, DOCX, and TXT files are supported.")
-
-# Determine aggregation strategy
-aggregation = "simple" if model_checkpoint in ["akdeniz27/xlm-roberta-base-turkish-ner", "xlm-roberta-large-finetuned-conll03-english", "asahi417/tner-xlm-roberta-base-ontonotes5"] else "first"
-
-st.subheader("Select Text Input Method")
-input_method = st.radio("", ('Write or Paste New Text', 'Upload File'))
-
-if input_method == "Write or Paste New Text":
-    input_text = st.text_area('Write or Paste Text Below', value="", height=128)
-else:
-    uploaded_file = st.file_uploader("Choose a file", type=["txt", "pdf", "docx"])
-    if uploaded_file is not None:
-        input_text = read_file(uploaded_file)
-        if input_text:
-            st.text_area("Extracted Text", input_text, height=128)
-    else:
-        input_text = ""
-
-@st.cache_resource
-def setModel(model_checkpoint, aggregation):
-    model = AutoModelForTokenClassification.from_pretrained(model_checkpoint)
-    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
-    return pipeline('ner', model=model, tokenizer=tokenizer, aggregation_strategy=aggregation)
-
-@st.cache_resource
-def entity_comb(output):
-    output_comb = []
-    for ind, entity in enumerate(output):
-        if ind == 0:
-            output_comb.append(entity)
-        elif output[ind]["start"] == output[ind-1]["end"] and output[ind]["entity_group"] == output[ind-1]["entity_group"]:
-            output_comb[-1]["word"] += output[ind]["word"]
-            output_comb[-1]["end"] = output[ind]["end"]
-        else:
-            output_comb.append(entity)
-    return output_comb
+# ... [Previous functions remain unchanged] ...
 
 def create_mask_dict(entities):
     mask_dict = {}
@@ -105,14 +21,13 @@ def create_mask_dict(entities):
                     entity_counters[entity['entity_group']] += 1
                 mask_dict[entity['word']] = f"{entity['entity_group']}_{entity_counters[entity['entity_group']]}"
     return mask_dict
-def create_masked_text(input_text, entities, mask_dict):
+
+def create_masked_text(input_text, mask_dict):
     masked_text = input_text
-    for entity in sorted(entities, key=lambda x: x['start'], reverse=True):
-        if entity['entity_group'] not in ['CARDINAL', 'EVENT']:
-            masked_text = masked_text[:entity['start']] + mask_dict[entity['word']] + masked_text[entity['end']:]
+    for word, mask in sorted(mask_dict.items(), key=lambda x: len(x[0]), reverse=True):
+        masked_text = re.sub(r'\b' + re.escape(word) + r'\b', mask, masked_text)
     return masked_text
 
-
 Run_Button = st.button("Run")
 
 if Run_Button and input_text:
@@ -134,47 +49,34 @@ if Run_Button and input_text:
                 entity['end'] += offset
         
         all_outputs.extend(output)
-        
     
     # Combine entities
-    
     output_comb = entity_comb(all_outputs)
     
     # Create mask dictionary
     mask_dict = create_mask_dict(output_comb)
 
-    masked_text = create_masked_text(input_text, output_comb, mask_dict)
-    
-    # Apply masking and add masked_word column
-    for entity in output_comb:
-        if entity['entity_group'] not in ['CARDINAL', 'EVENT']:
-            entity['masked_word'] = mask_dict.get(entity['word'], entity['word'])
-        else:
-            entity['masked_word'] = entity['word']
-    print("output_comb", output_comb)
-    #df = pd.DataFrame.from_dict(output_comb)
-    #cols_to_keep = ['word', 'entity_group', 'score', 'start', 'end']
-    #df_final = df[cols_to_keep].loc[:,~df.columns.duplicated()].copy()
+    # Create masked text
+    masked_text = create_masked_text(input_text, mask_dict)
     
-    #st.subheader("Recognized Entities")
-    #st.dataframe(df_final)
+    st.subheader("Masked Text")
+    st.text(masked_text)
 
-    
-    
-    # Spacy display logic with entity numbering
+    st.subheader("Masking Dictionary")
+    st.json(mask_dict)
+
+    # Create a DataFrame for display
+    df = pd.DataFrame([(word, mask) for word, mask in mask_dict.items()], columns=['Original', 'Masked'])
+    st.subheader("Masking Table")
+    st.dataframe(df)
+
+    # Optional: Display original text with highlights
+    st.subheader("Original Text with Highlights")
     spacy_display = {"ents": [], "text": input_text, "title": None}
     for entity in output_comb:
         if entity['entity_group'] not in ['CARDINAL', 'EVENT']:
-            label = f"{entity['entity_group']}_{mask_dict[entity['word']].split('_')[1]}"
-        else:
-            label = entity['entity_group']
-        spacy_display["ents"].append({"start": entity["start"], "end": entity["end"], "label": label})
+            label = mask_dict[entity['word']]
+            spacy_display["ents"].append({"start": entity["start"], "end": entity["end"], "label": label})
     
     html = spacy.displacy.render(spacy_display, style="ent", minify=True, manual=True)
-    st.write(html, unsafe_allow_html=True)
-
-    st.subheader("Masking Dictionary")
-    st.json(mask_dict)
-
-    st.subheader("Masked Text Preview")
-    st.text(masked_text)
+    st.write(html, unsafe_allow_html=True)