import gradio as gr
import requests
import json
import os
from screenshot import (
    before_prompt,
    prompt_to_generation,
    after_generation,
    js_save,
    js_load_script,
)

def inference(input_sentence, max_length, seed=42):
       parameters = {
            "max_new_tokens": max_length,
            "do_sample": False,
            "seed": seed,
            "early_stopping": False,
            "length_penalty": 0.0,
            "eos_token_id": None,
       }

    payload = {
        "inputs": input_sentence, 
        "parameters": parameters,
        "options" : {
            "use_cache": False
            } 
         }

    data = query(payload)

    if "error" in data:
        return (None, None, f"<span style='color:red'>ERROR: {data['error']} </span>")

    generation = data[0]["generated_text"].split(input_sentence, 1)[1]
    return (
        before_prompt
        + input_sentence
        + prompt_to_generation
        + generation
        + after_generation,
        data[0]["generated_text"],
        "",
    )


if __name__ == "__main__":
    demo = gr.Blocks()
    with demo:
        with gr.Row():
            gr.Markdown(value=description)
        with gr.Row():
            with gr.Column():
                text = gr.Textbox(
                    label="Input",
                    value=" ",  # should be set to " " when plugged into a real API
                )
                tokens = gr.Slider(1, 64, value=32, step=1, label="Tokens to generate")

                with gr.Row():
                    submit = gr.Button("Submit")
            with gr.Column():
                text_error = gr.Markdown(label="Log information")
                text_out = gr.Textbox(label="Output")
                display_out.set_event_trigger(
                    "load",
                    fn=None,
                    inputs=None,
                    outputs=None,
                    no_target=True,
                    js=js_load_script,
                )
        with gr.Row():
            gr.Examples(examples=examples, inputs=[text, tokens, sampling, sampling2])

        submit.click(
            inference,
            inputs=[text, tokens, sampling, sampling2],
            outputs=[display_out, text_out, text_error],
        )

    demo.launch()