from aitextgen import aitextgen
import gradio as gr
import os
from transformers import pipeline
from gradio import inputs
from gradio.inputs import Textbox
from gradio import outputs


ai=aitextgen(model='EleutherAI/gpt-neo-2.7B',to_gpu=False) # EleutherAI/gpt-neo-2.7B EleutherAI/gpt-neo-1.3B
# ai=aitextgen(model='EleutherAI/gpt-neo-1.3B',to_gpu=False)

def ai_text(Input):
  generated_text = ai.generate_one(max_length = 1000, prompt = Input, no_repeat_ngram_size = 3) #repetition_penalty = 1.9)
  #print(type(generated_text))
  return generated_text


title_ = "AI Long Content Generation"
description_ = "Converts short sentences into 1000 words"
output_text = gr.outputs.Textbox()
iface=gr.Interface(ai_text,"textbox", output_text, title=title_,description=description_)#.launch()
iface.launch()



#HF_TOKEN = os.environ.get("HF_TOKEN") 
#generator2 = gr.Interface.load("huggingface/EleutherAI/gpt-neo-2.7B", api_key=HF_TOKEN) # add api_key=HF_TOKEN to get over the quota error
#generator3 = gr.Interface.load("huggingface/EleutherAI/gpt-j-6B", api_key=HF_TOKEN)
#generator1 = gr.Interface.load("huggingface/gpt2-large", api_key=HF_TOKEN)
#gr.Parallel(generator1, generator2, generator3, inputs=gr.inputs.Textbox(lines=5, label="Enter a sentence to get another sentence."),
#            title=title, examples=examples).launch(share=False)