from gradio_pdf import PDF
from src.helper import *
import gradio as gr
from gradio_pdf import PDF

with gr.Blocks() as ner:
  gr.Markdown("# Sistem Ekstraksi Informasi Dokumen Putusan Hukum")
  # List Label
  keterangan_label = [
      ["VERN", "Nomor Putusan"],
      ["DEFN", "Nama Terdakwa"],
      ["CRIA", "Tindak Pidana"],
      ["ARTV", "Melanggar KUHP"],
      ["PENA", "Tuntutan Hukum"],
      ["PUNI", "Putusan Hukum"],
      ["TIMV", "Tanggal Putusan"],
      ["JUDP", "Hakim Ketua"],
      ["JUDG", "Hakim Anggota"],
      ["REGI", "Panitera"],
      ["PROS", "Penuntut Umum"],
      ["ADVO", "Pengacara"],
  ]
  gr.Markdown("## Penjelasan Label")
  gr.DataFrame(keterangan_label, headers=["Label", "Keterangan"], height=200)
  gr.Markdown("## Ekstraksi Entitas pada Potongan Kalimat")
  # Input Text
  with gr.Row():
    with gr.Column(scale=2):
      text = gr.Textbox(label="Text")
      model_text = gr.Dropdown(['IndoBERT (IndoLEM)', 'IndoBERT (IndoNLU)'], label='Model', value='IndoBERT (IndoLEM)', info='Pilih Model yang ingin digunakan *Default : IndoBERT (IndoLEM)')
      button_text = gr.Button(value="Predict", variant='primary')
      gr.ClearButton(text, value='Reset')
    with gr.Column(scale=3):
      output_text = gr.HighlightedText(label="Output Text")

    button_text.click(fn=text_extraction, inputs=[text, model_text], outputs=output_text, api_name="text")

  gr.Markdown("## Contoh Inputan Potongan Kalimat")
  gr.Examples(
    examples=[
        ["PUTUSAN . NOMOR : 187 / Pid . Sus / 2014 / PN . JKT . TIM . DEMI KEADILAN BERDASARKAN KETUHANAN YANG MAHA ESA . MENUNTUT : 1 Menyatakan terdakwa AGNES TRI AHADI Als AGNES telah terbukti secara sah dan meyakinkan bersalah melakukan tindak pidana Narkotika memiliki , menyimpan , menguasai , atau menyediakan Narkotika golongan I bukan tanaman sebagaimana didakwakan dalam dakwaan kedua yaitu melanggar ketentuan unsure pasal 112 ayat ( 1 ) UURI No . 35 tahun 2009 tentang Narkotika ;", "IndoBERT (IndoLEM)"],
        ["PUTUSAN . NOMOR : 187 / Pid . Sus / 2014 / PN . JKT . TIM", "IndoBERT (IndoNLU)"]
    ],
    inputs=[text, model_text],
    outputs=output_text,
    fn=text_extraction,
    )

  gr.Markdown("## Ekstraksi Entitas pada Dokumen Putusan Hukum")
  # Input PDF
  with gr.Row():
    with gr.Column(scale=2):
      doc = PDF(label="Document")
      model_pdf = gr.Dropdown(['IndoBERT (IndoLEM)', 'IndoBERT (IndoNLU)'], label='Model',value='IndoBERT (IndoLEM)', info='Pilih Model yang ingin digunakan *Default : IndoBERT (IndoLEM)')
      button_pdf = gr.Button(value="Extract", variant='primary')
      gr.ClearButton(doc, value="Reset")

    with gr.Column(scale=3):
      output_pdf = gr.Textbox(label="Output PDF")

  button_pdf.click(fn=pdf_extraction, inputs=[doc, model_pdf], outputs=output_pdf, api_name="pdf")

  gr.Examples(
    ["data/428_pid.b_2021_pn_jkt.brt_20240529091234.pdf",
     "data/1558_pid.b_2020_pn_jkt.brt_20240529091451.pdf",
     "data/329_pid.b_2023_pn_jkt.brt_20240529090837.pdf",
     "data/168_Pid.Sus_2023_PN_Bkl.pdf",
     "data/169_Pid.Sus_2023_PN_Bkl.pdf",
     "data/167_Pid.Sus_2023_PN_Bkl.pdf"],
    inputs=[doc],
    outputs=output_pdf,
    fn=pdf_extraction,
    )

if __name__ == "__main__":
  ner.launch()