import streamlit as st
from transformers import pipeline

# 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:
    st.json(base)
    st.json(mendo)