#GPT2 FINE TUNE

# library 
import gradio as gr
import tensorflow as tf  
from transformers import pipeline, AutoTokenizer, TFGPT2LMHeadModel


# function to run
def run_model(input_text, max_length):

    #GPT2 FT Transformer 
    model_gpt2 = TFGPT2LMHeadModel.from_pretrained("Sultannn/gpt2-ft-id-puisi")
    token_gpt2 = AutoTokenizer.from_pretrained("Sultannn/gpt2-ft-id-puisi")

    # preprocessing text input
    input_text = str(input_text)
    input_text = ' '.join(input_text.split()).upper() # hapus white space dan ubah judul kata ke huruf besar

    #encode input to vector
    text_prompt = token_gpt2.encode(input_text,add_special_tokens=True, return_tensors="tf") 

        
    #generate input
    summary_ids = model_gpt2.generate(text_prompt, # raw token 
                                    # param penting
                                    max_length = max_length + len(input_text), # max token to generate
                                    do_sample = True, # to sampling text 
                                                                        
                                    #length_penalty = length_penalty,  #Atur ke nilai <1.0 untuk menghasilkan urutan yang lebih pendek, ke nilai > 1.0 untuk menghasilkan urutan yang lebih panjang)    

                                    # pencarian modern          
                                    top_k = 50, # ambil top kata dengan probability tertinggi / kata yang paling mungkin
                                    top_p = 0.95, # ambil kata dan jumlah kan probalitiy nya sesuai yang di definisikan

                                    # no_repeat_ngram_size=2 , # agar tidak ada 2 gram/kata yang muncul dua kali:
                                    # temperature=1.0, # untuk mengatur probability next word 

                                    # trick / metode pencarian tradisional
                                    # num_beams=3, # Pencarian mengurangi risiko hilangnya urutan kata probabilitas tinggi yang tersembunyi dengan menjaga hipotesis yang paling mungkin pada setiap langkah waktu dan akhirnya memilih hipotesis yang memiliki probabilitas tertinggi secara keseluruhan 
                                    # early_stopping=True, # gunakan jika pakai num_beams > 1 
                                    # repetition_penalty=1.0, # mencegah kata pengulangan

                                    # jika ingin lebih dari 1 output
                                    # num_return_sequences=3,# num_return_sequences <= num_beams! jumlah balok skor tertinggi yang harus dikembalikan ? jumlah sample yang ingin dikeluarkan sesuai yang didefinisikan
                                    )                                                                       
                                        

    #decode output to text
    output = token_gpt2.decode(summary_ids[0], clean_up_tokenization_spaces=False) 

    return output # get output to str
                             
# end 

#example 
contoh = [["TAMPAN", 100],["TIDAK JELAS", 100]]
          
#judul
title = "indonesia Puisi Generator"

#deskripsi
description = "Demo for Puisi Generator ID. Models are GPT-2"           

#footer            
article = "<p style='text-align: center'><a href='https://github.com/sultanbst123/Hugging-Face-indo' target='_blank'><u>Untuk penjelasan lihat di repo ku</u> 😁</a></p>"

#run gradio
run = gr.Interface(
    fn=run_model,
    #input text
    inputs=[
        gr.inputs.Textbox(
            lines=2,
            placeholder="Ketik disini...", 
            label="Text",
        ),
        #fine tune

        #max length
        gr.inputs.Slider(
           minimum=100,
           maximum=150,  
           step=5,
           default=100, 
           label="Max Length(panjang maksimum urutan)",
       ),     
    ],
    #output text
    outputs=
    gr.outputs.Textbox(       
            label="Output text", 
    ),
    title=title,
    description=description,
    article=article,
    examples=contoh)
    
run.launch(debug = True)



## GOOD LUCK