import gradio as gr
import requests
import json
import os
from pathlib import Path

title = "🌸 BLOOM 🌸"
description = """Gradio Demo for using BLOOM with Spanish prompts. Heavily based on [Bloom demo](https://huggingface.co/spaces/huggingface/bloom_demo)
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 sentences of a webpage similar to the content you want to generate.
Start a paragraph as if YOU were writing a blog, webpage, math post, coding article and BLOOM will generate a coherent follow-up. Longer prompts usually give more interesting results.
Options:
- sampling: imaginative completions (may be not super accurate e.g. math/history)
- greedy: accurate completions (may be more boring or have repetitions)
"""

API_URL =  os.getenv("API_URL")
API_TOKEN = os.getenv("API_TOKEN")

examples = [
    ['Traduce español de España a español de Argentina\nEl coche es rojo - el auto es rojo\nEl ordenador es nuevo - la computadora es nueva\nel boligrafo es negro -', 16, "Sample"],
    ['Estos ejemplos quitan vocales de las palabras\nEjemplos:\nhola - hl\nmanzana - mnzn\npapas - pps\nalacran - lcrn\npapa -', 16, "Sample"],
    ["Un ejemplo de ecuación sencilla sería:\n4x = 40 ; en este caso el valor de x es", 16, "Greedy"],
    ["Si Pedro tiene 4 manzanas y María le quita 2, entonces a Pedro le quedan", 16, "Sample"],
    ["Esta es una conversación entre el modelo de lenguaje BLOOM y uno de sus creadores:\nCreador: Hola, BLOOM! ¿Tienes sentimientos?\nBLOOM:", 32, "Sample"],
    ["Había una vez un circo que alegraba siempre el", 32, "Sample"],
    ['''A continuación se clasifican reseñas de películas:\nComentario: "La película fue un horror"\nEtiqueta: Mala\n\nComentario: "La película me gustó mucho"\nEtiqueta: Buena\n\nComentario: "Es un despropósito de película"\nEtiqueta:''', 16, "Greedy"],
    ['''# La siguiente función hace un petición a la API y devuelve la respuesta en formato JSON\ndef query(payload, model_id, api_token):\n\theaders = {"Authorization": f"Bearer {api_token}"}\n\tAPI_URL = f"https://api-inference.huggingface.co/models/{model_id}"\n\tresponse =''',32, "Sample"],
    ['''Ingredientes de la paella:\n\nArroz bomba - 1500 g\nPollo de corral - 1\nConejo - 0.5 kg\nJudía verde plana''', 32, "Sample"],
    ['''En Barcelona podemos visitiar los siguientes edificios:\n\n- La Sagrada Familia\n- Las Ramblas''', 32, "Sample"],
    ['''Completa la siguientes series numéricas:\n\n2, 4, 6, 8\n1, 3, 5, 7\n3, 5, 7,''', 32, "Sample"] 
]

def query(payload):
    print(payload)
    headers = {"Authorization": f"Bearer {API_TOKEN}"}
    response = requests.request("POST", API_URL, headers=headers, 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"], label="Decoding", default="Sample")
    ], 
    ["text"],
    examples=examples,
    # article=article,
    cache_examples=False,
    title=title,
    description=description
).launch()