Update app.py

app.py

@@ -1,9 +1,78 @@
 import streamlit as st
 from transformers import pipeline
 
-pipe = pipeline('sentiment-analysis')
+# Load the pre-trained NER model
+model = pipeline("ner", model="/ner-app/mendobert/", tokenizer="indolem/indobert-base-uncased")
+basemodel = pipeline("ner", model="/ner-app/base-model/", tokenizer="indolem/indobert-base-uncased")
+
+ner_results = model(text)
+ner_results2 = basemodel(text)
+
+
+# MendoBERT
+
+formatted_results = []
+for result in ner_results:
+  end = result["start"]+len(result["word"].replace("##", ""))  
+  
+  if result["word"].startswith("##"):
+    formatted_results[-1]["end"] = end
+    formatted_results[-1]["word"]+= result["word"].replace("##", "")
+  else:
+    formatted_results.append({
+            'start': result["start"], 
+            'end': end,
+            'entity': result["entity"],
+            'index': result["index"],
+            'score': result["score"],
+            'word': result["word"]})
+    
+for result in formatted_results:
+    if result["entity"].startswith("LABEL_0"):
+        result["entity"] = "O"
+    elif result["entity"].startswith("LABEL_1"):
+        result["entity"] = "B"
+    elif result["entity"].startswith("LABEL_2"):
+        result["entity"] = "I"
+
+mendo =[]
+for result in formatted_results:
+    if not result["entity"].startswith("O"):
+        mendo.append(f"""Entity: {result["entity"]}, Start:{result["start"]}, End:{result["end"]}, word:{text[result["start"]:result["end"]]}""")        
+
+# Base Model     
+
+formatted_results = []
+for result in ner_results2:
+  end = result["start"]+len(result["word"].replace("##", ""))  
+  
+  if result["word"].startswith("##"):
+    formatted_results[-1]["end"] = end
+    formatted_results[-1]["word"]+= result["word"].replace("##", "")
+  else:
+    formatted_results.append({
+            'start': result["start"], 
+            'end': end,
+            'entity': result["entity"],
+            'index': result["index"],
+            'score': result["score"],
+            'word': result["word"]})
+    
+for result in formatted_results:
+    if result["entity"].startswith("LABEL_0"):
+        result["entity"] = "O"
+    elif result["entity"].startswith("LABEL_1"):
+        result["entity"] = "B"
+    elif result["entity"].startswith("LABEL_2"):
+        result["entity"] = "I"
+
+base=[]        
+for result in formatted_results:
+    if not result["entity"].startswith("O"):
+        base.append(f"""Entity: {result["entity"]}, Start:{result["start"]}, End:{result["end"]}, word:{text[result["start"]:result["end"]]}""")
+
 text = st.text_area('enter some text: ')
 
 if text:
-    out = pipe(text)
-    st.json(out)
+    st.json(base)
+    st.json(mendo)