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.
Tips: Do NOT talk to BLOOM as an entity, it's not a chatbot but a webpage/blog/article completion model. For the best results: MIMIC a few words of a webpage similar to the content you want to generate. Start a sentence as if YOU were writing a blog, webpage, math post, coding article and BLOOPM will generate a coherent follow-up.

Generation simple selector:
- sampling: imaginative completions (may be not super accurate e.g. math/history)
- greedy: accurate completions (may be more boring or have repetitions)
"""

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:', 32, "Sample"],
    ["Pour déguster un ortolan, il faut tout d'abord", 32, "Sample"],
    ["Question: If I put cheese into the fridge, will it melt?\nAnswer:", 32, "Sample"],
    ["Math exercise - answers:\n34+10=44\n54+20=", 16, "Greedy"],
    ["Question: Where does the Greek Goddess Persephone spend half of the year when she is not with her mother?\nAnswer:", 24, "Greedy"],
    ["spelling test answers.\nWhat are the letters in « language »?\nAnswer: l-a-n-g-u-a-g-e\nWhat are the letters in « Romanian »?\nAnswer:", 24, "Greedy"],
]

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, sample_or_greedy, seed=42):
    if sample_or_greedy == "Sample":
        parameters = {"max_new_tokens": max_length,
                      "top_p": 0.9,
                      "do_sample": True,
                      "seed": seed,
                      "early_stopping": False,
                      "length_penalty": 0.0,
                      "eos_token_id": None}
    else:
        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}

    data = query(
        payload
    )
    print(data)
    return data[0]["generated_text"]


gr.Interface(
    inference, 
    [
        gr.inputs.Textbox(label="Input"),
        gr.inputs.Slider(1, 64, default=32, step=1, label="Tokens to generate"),
        gr.inputs.Radio(["Sample", "Greedy"]),
    ], 
    gr.outputs.Textbox(label="Output"),
    examples=examples,
    # article=article,
    title=title,
    description=description,
    flagging_callback=hf_writer,
    allow_flagging=True,
).launch()