import gradio as gr
from transformers import pipeline
pipeline = pipeline("text-generation", model="not-lain/PyGPT")

def format_input(instruction,inp):
  prefix = f"Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n"
  txt = prefix + f"### Instruction:\n{instruction}"+ f"\n\n### Input:{inp}"+"\n\n### Output:\n"
  return txt

def process_markdown(out):
  mark = f"""```py
{out}
```"""
  return mark

def generate_text(length,instruction,inp):
    if instruction == None  : 
       instruction = ""
    if inp == None  : 
       inp = ""
    txt = format_input(instruction,inp)
    out =  pipeline(txt, max_length=len(txt)+length)[0]["generated_text"]
    out = out.split("Output:\n")[1]
    mark = process_markdown(out)
    return mark


MARKDOWN_TEXT = """
# PyGPT Text Generation Demo
this is a demo using the [PyGPT model](https://huggingface.co/not-lain/PyGPT) to generate text based on an input instruction and input text.
the model is based on the [GPT-2 model](https://huggingface.co/gpt2) and finetuned on the [python_code_instructions_18k_alpaca](https://huggingface.co/datasets/iamtarun/python_code_instructions_18k_alpaca) dataset.
"""



with gr.Blocks() as iface:
    gr.Markdown(MARKDOWN_TEXT)
    length = gr.Slider(1, 100, 50, label="Max Length")
    instruction = gr.Text(label= "instruction")
    inp = gr.Text(label="input")
    out = gr.Markdown(label="output")
    submit = gr.Button("submit")
    submit.click(generate_text,inputs=[length,instruction,inp],outputs=out)
    gr.Examples([
        [50,"Create a function to calculate the sum of a sequence of integers.","[1, 2, 3, 4, 5]"],
        [50,"Generate a Python code for crawling a website for a specific type of data.","website: www.example.com data to crawl: phone numbers"]],
                inputs = [length,instruction,inp],
                outputs= [out],
                fn=generate_text,
                cache_examples=True)

    

iface.launch(debug=True)