Upload app.py

app.py

@@ -1,13 +1,15 @@
 import gradio as gr
+import random
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-model = AutoModelForCausalLM.from_pretrained("aronvandepol/K-GPT125M", use_auth_token="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")
+
+model = AutoModelForCausalLM.from_pretrained("aronvandepol/K-GPT125M", use_auth_token=os.environ['TOKEN'])
 
 tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-125M")
 
 model = AutoModelForCausalLM.from_pretrained("K-GPT_NEO125M")
 
-def get_gen(sample: str, returns: int=1, length: int=50):
+def get_gen(sample: str, length: int=30, beam: int =1):
     input_ids = tokenizer(sample, return_tensors='pt').input_ids
     output = model.generate(input_ids, 
         max_length = length, 
@@ -15,15 +17,45 @@ def get_gen(sample: str, returns: int=1, length: int=50):
         no_repeat_ngram_size  = 3,
         early_stopping = True,
         do_sample=True,
-        num_return_sequences = returns,
+        num_return_sequences = 10,
         pad_token_id=tokenizer.eos_token_id
     )
     generation = tokenizer.batch_decode(output, skip_special_tokens=True)
-    return generation[0]
+    gen_text = generation[beam-1]
+    return gen_text
+
 
-with gr.Blocks() as demo:
-    textbox = gr.Textbox(placeholder="Type here and press enter...", lines=4)
-    btn = gr.Button("Generate")
-    btn.click(get_gen, textbox, textbox)
+with gr.Blocks() as app:
+    gr.Markdown("<h1><center>K-GPT_NEO</center></h1>")
+    gr.Markdown(
+        """
+        <justify>
+        Interact with the K-GPT_NEO model and generate kpop song texts! By entering a few words into the <i>input</i> prompt and press generate to get the most probable sentence. If you want to see some less probable results press the <i>I'm feeling lucky</i> button.
+        </justify>
+        """
+        )
+    beam=gr.Number(value=1, visible=False, precision=0)
+    with gr.Row():
+        length = gr.Slider(0, 100, step=5, label="Max generated words", value=30)
+    with gr.Group():
+        txt1 = gr.Textbox(label = "Input", placeholder="Type here and press enter...", lines=4)
+        txt2 = gr.Textbox(label = "Output", placeholder="Generated sentence will appear here", lines=4, interactive=False)
+    with gr.Row():
+        btn = gr.Button("Generate most probable")
+        rnd = gr.Button("Feeling Lucky!")
+    btn.click(fn=get_gen, inputs=[txt1, length, beam], outputs=txt2)
+    rnd.click(fn=get_gen, inputs=[txt1, length, gr.Number(value=random.randint(2, 10), visible=False, precision=0)], outputs=txt2)
+    gr.Examples(examples=[['I miss you'], ['My Love has not faded yet'], ['Dancing the stars away']], inputs=txt1, outputs=txt2, fn=get_gen, cache_examples=True)
+    gr.Markdown(
+        """
+        <br>
+        <br>
+        <br>
+        <justify>
+        K-GPT_NEO is based on the GPT-Neo text generation model developed by Eleuther AI. 
+        This architecture was fine-tuned on 2000 English translations of K-pop songs ranging from BTS and BLACKPINK, to TWICE and ZONE. For more information on the training and data, please visit: 
+        </justify>
+        """
+        )
 
-demo.launch(share=False)
+app.launch()