app.py

import streamlit as st
import itertools
from typing import Dict, Union

from nltk import sent_tokenize
import nltk
nltk.download('punkt')
import torch
from transformers import(
    AutoModelForSeq2SeqLM, 
    AutoTokenizer
)

class QGPipeline:

    def __init__(
        self
    ):
      
        self.model = AutoModelForSeq2SeqLM.from_pretrained("muchad/idt5-qa-qg")
        self.tokenizer = AutoTokenizer.from_pretrained("muchad/idt5-qa-qg")
        self.qg_format = "highlight"
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model.to(self.device)
        self.ans_model = self.model
        self.ans_tokenizer = self.tokenizer
        assert self.model.__class__.__name__ in ["T5ForConditionalGeneration"]
        self.model_type = "t5"


    def __call__(self, inputs: str):
        inputs = " ".join(inputs.split())
        sents, answers = self._extract_answers(inputs)
        flat_answers = list(itertools.chain(*answers))

        if len(flat_answers) == 0:
          return []

        qg_examples = self._prepare_inputs_for_qg_from_answers_hl(sents, answers)        
        qg_inputs = [example['source_text'] for example in qg_examples]
        questions = self._generate_questions(qg_inputs)
        output = [{'answer': example['answer'], 'question': que} for example, que in zip(qg_examples, questions)]
        return output
    
    def _generate_questions(self, inputs):
        inputs = self._tokenize(inputs, padding=True, truncation=True)
        
        outs = self.model.generate(
            input_ids=inputs['input_ids'].to(self.device), 
            attention_mask=inputs['attention_mask'].to(self.device), 
            max_length=80,
            num_beams=4,
        )
        
        questions = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in outs]
        return questions
    
    def _extract_answers(self, context):
        sents, inputs = self._prepare_inputs_for_ans_extraction(context)
        
        inputs = self._tokenize(inputs, padding=True, truncation=True)

        outs = self.ans_model.generate(
            input_ids=inputs['input_ids'].to(self.device), 
            attention_mask=inputs['attention_mask'].to(self.device), 
            max_length=80,
        )
        
        dec = [self.ans_tokenizer.decode(ids, skip_special_tokens=True) for ids in outs] 
        answers = [item.split('<sep>') for item in dec]
        answers = [i[:-1] for i in answers]
        return sents, answers
    
    def _tokenize(self,
        inputs,
        padding=True,
        truncation=True,
        add_special_tokens=True, 
        max_length=512
    ):
        inputs = self.tokenizer.batch_encode_plus(
            inputs, 
            max_length=max_length,
            add_special_tokens=add_special_tokens,
            truncation=truncation,
            padding="max_length" if padding else False,
            pad_to_max_length=padding,
            return_tensors="pt"
        )
        return inputs
    
    def _prepare_inputs_for_ans_extraction(self, text):
        sents = sent_tokenize(text)

        inputs = []
        for i in range(len(sents)):
            source_text = "extract answers:"
            for j, sent in enumerate(sents):
                if i == j:
                    sent = "<hl> %s <hl>" % sent
                source_text = "%s %s" % (source_text, sent)
                source_text = source_text.strip()
            
            source_text = source_text + " </s>"
            inputs.append(source_text)
        return sents, inputs
    
    def _prepare_inputs_for_qg_from_answers_hl(self, sents, answers):
        inputs = []
        for i, answer in enumerate(answers):
            if len(answer) == 0: continue
            for answer_text in answer:
                sent = sents[i]
                sents_copy = sents[:]
                
                answer_text = answer_text.strip()
                try:
                  ans_start_idx = sent.index(answer_text)
                  
                  sent = f"{sent[:ans_start_idx]} <hl> {answer_text} <hl> {sent[ans_start_idx + len(answer_text): ]}"
                  sents_copy[i] = sent
                  
                  source_text = " ".join(sents_copy)
                  source_text = f"generate question: {source_text}" 
                  if self.model_type == "t5":
                      source_text = source_text + " </s>"
                except:
                  continue
                
                inputs.append({"answer": answer_text, "source_text": source_text})
        
        return inputs
    
class TaskPipeline(QGPipeline):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
    
    def __call__(self, inputs: Union[Dict, str]):
        return super().__call__(inputs)
       
def pipeline():
    task = TaskPipeline
    return task()
  
@st.cache(ttl=24*3600,allow_output_mutation=True)
def pipeline():    
    task = TaskPipeline    
    return task()

qg = pipeline()
st.title("Indonesian Question Generation")
st.write("Indonesian Question Generation System using [idT5](https://huggingface.co/muchad/idt5-base)")
default_context = "Kapitan Pattimura adalah pahlawan dari Maluku. Beliau lahir pada tanggal 8 Juni 1783 dan meninggal pada tanggal 16 Desember 1817."
option = st.selectbox(
    'Example:',
    (default_context, 
     'Dua orang pengembara berjalan di sepanjang jalan yang berdebu dan tandus di hari yang sangat panas. Tidak lama kemudian, mereka menemukan sebuah pohon besar.', 
     'Raja Purnawarman mulai memerintah Kerajaan Tarumanegara pada tahun 395 M. Pada masa pemerintahannya, ia selalu berjuang untuk rakyatnya. Ia membangun saluran air dan memberantas perompak.\r\r\nRaja Purnawarman sangat memperhatikan kesejahteraan rakyatnya. la memperbaiki aliran Sungai Gangga di daerah Cirebon. Dua tahun kemudian, ia juga memperbaiki dan memperindah alur Sungai Cupu sehingga air bisa mengalir ke seluruh kerajaan. Para petani senang karena ladang mereka mendapat air dari aliran sungai sehingga menjadi subur. Ladang para petani tidak kekeringan pada musim kemarau.',
     'Pada masa tersebut, buku peraga Ini Budi laris manis tidak hanya di Jakarta, Bengkulu, Palembang, dan kota lain melainkan hampir di seluruh kota di Indonesia. Namun Siti Rahmani Rauf mengakui bahwa buku paket Ini Budi sebenarnya telah ada namun buku peraganya yang tidak ada, oleh karena itu di buatlah buku peraga Ini Budi, yang tidak lain adalah pembelajaran yang pernah kita jumpai ketika belajar di bangku kelas 1 SD. Karena buku peraga itulah yang membuat kalimat Ini Budi menjadi sangat populer hingga saat ini untuk di pelajari oleh siswa kelas 1 SD.',
     'Rabiul Awal atau bulan ketiga dalam kalender Hijriyah adalah bulan yang dimuliakan umat Islam. Sebab, di bulan ini Nabi Muhammad SAW lahir ke dunia sebagai manusia yang diutus oleh Allah SWT untuk menyampaikan ajaran-Nya.\r\r\nJauh sebelum kelahiran Rasulullah, Allah sudah mengabarkan akan kehadiran Nabi akhir zaman. Kedatangan Rasulullah telah disebut-sebut dalam kitab sebelum Alquran, yakni dalam kitab Taurat dan Injil. Sehingga, para rabi Yahudi dan pendeta Nasrani telah mengenal Rasulullah dari gambaran tentang sifat-sifatnya.\r\r\nDikutip dari Republika Online.',
     'Pemerintah berencana menaikkan tiket Candi Borobudur menjadi Rp 750.000/orang. Kenaikan harga tiket ini bertujuan untuk membatasi jumlah pengunjung sehingga kondisi candi sebagai cagar budaya tetap ini tetap lestari. Namun, harga tiket yang mencapai Rp 750.000 per orang itu bukanlah tiket masuk, melainkan tiket naik ke atas Candi Borobudur. Direktur Utama PT Aviasi Pariwisata Indonesia (Persero)/InJourney, Dony Oskaria menyebut tiket masuk Candi Borobudur tak berubah.'))
context_in = st.text_area('Context:', option, height=200)
if st.button('Generate Question'):
    if context_in: 
        questions = qg(context_in)
        re = ""
        for i, q in enumerate(questions):
            re += "Q : %s" % q['question'] + "  \n" + "A : %s" % q['answer'] + "  \n\n" 
        st.write(re)
    else:
        st.write("Please check your context")