import gradio as gr
from transformers import pipeline
from gtts import gTTS
import torch

# GPU control
device = 0 if torch.cuda.is_available() else -1

# gpt2-large model text generation
def text_gen(prompt):
  generator = pipeline("text-generation", model='gpt2-medium')
  #generator = pipeline("text-generation", model='distilgpt2')
  #res_base = generator(prompt, max_new_tokens=100, num_return_sequences=1)
  #res_p = generator(prompt, max_new_tokens=100, num_return_sequences=1, do_sample=True, top_p=0.095)
  #res_k = generator(prompt, max_new_tokens=120, num_return_sequences=1, do_sample=True, top_k=4)
  #res_c = generator(prompt, max_new_tokens=100, num_return_sequences=1, penalty_alpha=0.6, top_k=4)
  #<span style="font-family:Papyrus; font-size:2em;">👽 0xZee</span>
  res = generator(prompt, max_lenght=120, num_return_sequences=1, device=device)
  res = res[0]['generated_text']
  # text_to_audio
  audio = gTTS(text=res, lang='en', slow=False)
  audio.save("res.wav")
  return(res , "res.wav")

def complete_with_gpt(text):
  complete = pipeline("text-generation", model='openai-gpt')
  res = complete(text[-50:], max_new_tokens=100, num_return_sequences=1, device=device)
  return res[0]['generated_text']


##
with gr.Blocks() as demo:
    with gr.Row():
      gr.Markdown("""
                  > 👽 0xZee
                  # 📝 NLP : GPT Text Generation & Text-to-Speech Demo
                  > Demo #1 : **gtp2** Model (finetuned) and **gTTS** for text to audio
                  *Enter a text with context : Will genrate text and display audio.*
                  """)
#
    with gr.Row():
        with gr.Column():   
          prompt = gr.Textbox(label="Text Generation", placeholder="Text and context here")
          gr.Examples(examples=["The Ai can help the world with", "The Moon's orbit around Earth has" , "John is a cowboy, suddenly he"], inputs=[prompt])
          btn = gr.Button(value="Go GPT")
#
        with gr.Column():   
          audio_out = gr.Audio()
          text_out = gr.Textbox(label="tuned mode")
          btn.click(fn=text_gen, inputs=prompt, outputs=[text_out,audio_out])
#
    with gr.Row():
      gr.Markdown("""
                  #
                  # 🎲 NLP : GPT Text Auto Complete
                  > Demo #2 : Auto-Complete : **openai-gpt** Model 
                  *Enter a text : will complete the text provided.*
                  """)
#
    with gr.Row():
          text = gr.Textbox(label="Text Complete", placeholder="Text here", lines=4)
    with gr.Row():
          gr.Examples(examples=["The Ai can help the world with", "tell me a story about" , "John is a cowboy, suddenly he is"], inputs=[text])     
    with gr.Row():          
          btn = gr.Button(value="Go & Complete")
          btn.click(fn=complete_with_gpt, inputs=text, outputs=text, queue=False)



if __name__ == "__main__":
  demo.launch(debug=True)