import gradio as gr
import requests
import os 

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

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


def sql_generate(prompt, input_prompt_sql ): 

  print(f"*****Inside SQL_generate - Prompt is :{prompt}")
  print(f"length of input_prompt_sql is {len(input_prompt_sql)}")
  print(f"length of prompt is {len(prompt)}")
  if len(prompt) == 0:
    prompt = input_prompt_sql 
    
  json_ = {"inputs": prompt,
            "parameters":
            {
            "top_p": 0.9,
          "temperature": 1.1,
          "max_new_tokens": 64,
          "return_full_text": False,
          }, 
          "options": 
          {"use_cache": True,
          "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']
  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 


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():
  
    example_prompt = gr.Radio( [
    "Instruction: Given an input question, respond with syntactically correct PostgreSQL\nInput: How many users signed up in the past month?\nPostgreSQL query: ", 
    "Instruction: Given an input question, respond with syntactically correct PostgreSQL\nInput: Create a query that displays empfname, emplname, deptid, deptname, location from employee table. Results should be in the ascending order based on the empfname and location.\nPostgreSQL query: ", 
    "Instruction: Given an input question, respond with syntactically correct PostgreSQL. Only use table called 'employees'.\nInput: What is the total salary paid to all the employees?\nPostgreSQL query: ",
    "Instruction: Given an input question, respond with syntactically correct PostgreSQL. Only use table called 'employees'.\nInput: List names of all the employees whose name end with 'r'.\nPostgreSQL query: ",
    "Instruction: Given an input question, respond with syntactically correct PostgreSQL. Only use table called 'employees'.\nInput: What are the number of employees in each department?\nPostgreSQL query: ",
    "Instruction: Given an input question, respond with syntactically correct PostgreSQL. Only use table called 'employees'.\nInput: Select names of all theemployees who have third character in their name as 't'.\nPostgreSQL query: ",
    "Instruction: Given an input question, respond with syntactically correct PostgreSQL. Only use table called 'employees'.\nInput: Select names of all the employees who are working under 'Peter'.\nPostgreSQL query: ",  ], label= "Choose a sample Prompt")
  
    #with gr.Column:  
    input_prompt_sql = gr.Textbox(label="Or 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=[example_prompt, input_prompt_sql], outputs=generated_txt) 

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

# import gradio as gr
# gr.Interface.load("models/bigscience/bloom").launch()