import gradio as gr
import re
import requests
import json
import os

title = "BLOOM"
description = "Gradio Demo for BLOOM. To use it, simply add your text, or click one of the examples to load them. Read more at the links below."

API_URL = "https://hfbloom.ngrok.io/generate"
HF_API_TOKEN = os.getenv("HF_API_TOKEN")

hf_writer = gr.HuggingFaceDatasetSaver(HF_API_TOKEN, "huggingface/bloom_internal_prompts", organization="huggingface")


examples = [
    ['A "whatpu" is a small, furry animal native to Tanzania. An example of a sentence that uses the word whatpu is: We were traveling in Africa and we saw these very cute whatpus. To do a "farduddle" means to jump up and down really fast. An example of a sentence that uses the word farduddle is:', 8, 0, 0, 0.1, 0, 0.9, False],
    ['def quicksort(l):', 8, 0, 0, 0.1, 0, 0.9, False],
    ['Q: ¿Cómo te llamas? A: What is your name? Q: ¿Qué edad tienes? A: How old are you? Q: ¿Dónde vives?', 8, 0, 0, 0.1, 0, 0.9, False]
]

def safe_text(text):
    text = text.replace('%', '\\%25')
    text = text.replace('#', '\\%23')
    text = text.replace('+', '\\%2B')
    text = text.replace('*', '\\%2A')
    text = text.replace('&', '\\%26')
    text = re.sub(r"([$_*\[\]()~`>\#\+\-=|\.!{}])", r"\\\1", text)
    return f"<pre>{text}</pre>"


def query(payload):
    print(payload)
    response = requests.request("POST", API_URL, json=payload)
    print(response)
    return json.loads(response.content.decode("utf-8"))
    
def inference(input_sentence, max_length, no_repeat_ngram_size, num_beams, temperature,top_k, top_p, greedy_decoding, seed=42):
    top_k = None if top_k == 0 else top_k
    do_sample = False if num_beams > 0 else not greedy_decoding
    num_beams = None if (greedy_decoding or num_beams == 0) else num_beams
    no_repeat_ngram_size = None if num_beams is None else no_repeat_ngram_size
    top_p = None if num_beams else top_p

    early_stopping = None if num_beams is None else num_beams > 0

    payload = {"inputs": input_sentence,
               "parameters": {"max_new_tokens": max_length, "top_k": top_k, "top_p": top_p, "temperature": temperature,
                              "do_sample": do_sample, "seed": seed, "early_stopping":early_stopping, "no_repeat_ngram_size":no_repeat_ngram_size, "num_beams":num_beams}}
    data = query(
        payload
    )
    print(data)
    return data[0]["generated_text"]


gr.Interface(
    inference, 
    [
        gr.inputs.Textbox(label="Input"),
        gr.inputs.Slider(1, 64, default=8, label="Tokens to generate"),
        gr.inputs.Slider(1, 10, default=2, step=1, label="No repeat N gram"),
        gr.inputs.Slider(0, 10, default=5, step=1, label="Num beams"),
        gr.inputs.Slider(0.0, 1.0, default=0.1, step=0.05, label="Temperature"),
        gr.inputs.Slider(0, 64, default=0, step=1, label="Top K"),
        gr.inputs.Slider(0.0, 10, default=0.9, step=0.05, label="Top P"),
        gr.inputs.Checkbox(False, label="Greedy decoding"),
    ], 
    gr.outputs.Textbox(label="Output"),
    examples=examples,
    # article=article,
    title=title,
    description=description,
    flagging_callback=hf_writer,
    allow_flagging=True,
).launch()