import gradio as gr
import requests
import os 

##Bloom
API_URL = "https://api-inference.huggingface.co/models/bigscience/bloom"
HF_TOKEN = os.environ["HF_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(input_prompt_sql) == 0:
    prompt = "Instruction: Given an input question, respond with syntactically correct PostgreSQL\nInput: " +input_prompt_sql #+ "\nPostgreSQL query: "
    
  json_ = {"inputs": prompt,
            "parameters":
            {
            "top_p": 0.9,
          "temperature": 1.1,
          "max_new_tokens": 250,
          "return_full_text": False,
          }, 
          "options": 
          {"use_cache": True,
          "wait_for_model": True,
          },}
  response = requests.post(API_URL, headers=headers, json=json_)
  print(f"Response  is : {response}")
  output = response.json()
  print(f"output is : {output}") #{output}")
  output_tmp = output[0]['generated_text']
  print(f"output_tmp is: {output_tmp}")
  solution = output_tmp.split("\nQ:")[0]   #output[0]['generated_text'].split("Q:")[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>Zero Shot 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\nThis Space is created by [Yuvraj Sharma](https://twitter.com/yvrjsharma) for EuroPython 2022 Demo.\nThe Prompt length is limited at the API end right now, thus there is a certain limitation in testing Bloom's capability thoroughly.\n\nThis 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, for the same prompt sometimes the Space runs just fine, while the other times it runs into error showing Response code [500] or *A very long prompt, temporarily not accepting these* message in the logs.\nStill iterating over the app, might be able to improve it further soon.."""
        )
  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 tables 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 tables 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 tables 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")
  
  
    input_prompt_sql = gr.Textbox(label="Or Write text following the example pattern given below, to get SQL commands...", value="Only use table called 'department'.\nInput: Select names of all the departments in descending alphabetical order of department names.\nPostgreSQL query: ", lines=5)
    
  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)