import gradio as gr
import requests
import os

##Bloom
API_URL = "https://api-inference.huggingface.co/models/bigscience/bloom"

HF_TOKEN = "Token"
headers = {"Authorization": f"Bearer {HF_TOKEN}"}


def sql_generate(prompt):
    json_ = {"inputs": prompt,
             "parameters":
                 {
                     "top_p": 0.9,
                     "temperature": 1.1,
                     "max_new_tokens": 64,
                     "return_full_text": True,
                 },
             "options":
                 {"use_cache": False,
                  "wait_for_model": True,
                  }, }
    response = requests.post(API_URL, json=json_)
    print(f"Response  is : {response}")
    output = response.json()
    print(f"output is : {output}")
    output_tmp = output[0]['generated_text']
    output_tmp= output_tmp.replace(prompt,'')
    print(f"output_tmp is: {output_tmp}")
    solution = output_tmp.split("\nQ:")[0]
    print(f"Final response after splits is: {solution}")
    if '\nOutput:' in solution:
        final_solution = solution.split("\nOutput:")[0]
        print(f"Response after removing output is: {final_solution}")
    elif '\n\n' in solution:
        final_solution = solution.split("\n\n")[0]
        print(f"Response after removing new line entries is: {final_solution}")
    else:
        final_solution = solution
    # return final_solution
    return final_solution


demo = gr.Blocks()

with demo:
    gr.Markdown("<h1><center>Text to SQL by Bloom</center></h1>")
    gr.Markdown(
        """[BigScienceW Bloom](https://twitter.com/BigscienceW) \n\n Large language models have demonstrated a capability of Zero-Shot SQL generation. Some might say — You can get good results out of LLMs if you know how to speak to them. This space is an attempt at inspecting this behavior/capability in the new HuggingFace BigScienceW [Bloom](https://huggingface.co/bigscience/bloom) model.\n\nThe Prompt length is limited at the API end right now, thus there is a certain limitation in testing Bloom's capability thoroughly.This Space might sometime fail due to inference queue being full and logs would end up showing error as *'queue full, try again later'*, in such cases please try again after few minutes. Please note that, longer prompts might not work as well and the Space could error out with Response code [500] or *'A very long prompt, temporarily not accepting these'* message in the logs. Still iterating over the app, might be able to improve it further soon.. \n\nThis Space is created by Yuvraj Sharma for Gradio EuroPython 2022 Demo and changed with [Ayoub Nikravan](https://www.linkedin.com/in/ayyoub-nikravan-862b1770/) ."""
    )
    with gr.Row():
        input_prompt_sql = gr.Textbox(
            label="Write text following the example pattern given below, to get SQL commands...",
            value="Instruction: Given an input question, respond with syntactically correct PostgreSQL. Use table called 'department'.\nInput: Select names of all the departments in their descending alphabetical order.\nPostgreSQL query: ",
            lines=6)

    with gr.Row():
        generated_txt = gr.Textbox(lines=3)

    b1 = gr.Button("Generate SQL")
    b1.click(sql_generate, inputs=[input_prompt_sql], outputs=generated_txt)

demo.launch(enable_queue=True, debug=True)