from flair.models import SequenceTagger
from flair.data import Sentence
import gradio as gr


# [{'entity': 'I-LOC',
#   'score': 0.9988978,
#   'index': 2,
#   'word': 'Chicago',
#   'start': 5,
#   'end': 12},
#  {'entity': 'I-MISC',
#   'score': 0.9958592,
#   'index': 5,
#   'word': 'Pakistani',
#   'start': 22,
#   'end': 31}]


def predict(text):
   model = SequenceTagger.load("yudiwbs/absa_elektronik_aspek_v2")
   sentence = Sentence(text)
   # predict
   model.predict(sentence)

   str = ""
   out = []
   for idx,entity in enumerate(sentence.get_spans('aspek')):
      ent = {}
      ent['entity'] = entity.tag
      ent['score'] = entity.score
      ent['index'] = idx
      ent['word'] = entity.text
      ent['start'] = entity.start_position
      ent['end'] = entity.end_position
      out.append(ent)


   return {"text":text, "entities":out} 


#predict("Sesuai diskusi. Barang mulus walau ada hairline scratch di casing depan. Battery masih awet. Dapat bonus jg. Thanks gan.")

examples = [
    "Sesuai diskusi. Barang mulus walau ada hairline scratch di casing depan. Battery masih awet. Dapat bonus keyboard Jepang jg. Thanks gan.",
    "product.. Ram 8GB HDD 320 sesuai pesanan. body luar dalam mulus, gores kecil pemakaian wajar.. keyboard empuk. display no dot Pixel. wifi. usb. sound. bluetooth. charger baterai ok awet 😁😁 over all sangat recommended mantaaaap jiwaaaaa",
    "overall, laptopnya compact, keyboard berfungsi semua, LED msh bagus, body bersih, charger berfungsi dgn baik, wifi jg connect, dan sudah terinstall win x. recommended" 
]


demo = gr.Interface(
  predict,
  gr.Textbox(placeholder="Enter sentence here..."),
  gr.HighlightedText(label="Aspek",color_map={"KEY": "red","BOD": "yellow", "POW": "green"}),
  examples = examples,
  theme=gr.themes.Default()
  #inputs='text',
  #outputs='text'
)

demo.launch()