import gradio as gr
import requests

# GPT-J-6B API
API_URL = "https://api-inference.huggingface.co/models/EleutherAI/gpt-j-6B"
headers = {"Authorization": "Bearer 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"}
prompt = """Customer: Hi, this is M. Davenport, how may I direct your call?
Agent: Thankyou, today I seek some Wellness and Mindfulness advice.
Customer: Great!  I've been searching for good solutions to enhance memory and health. 
Agent: Let me share some of the resources with you including mnemonics, agents, nutrition, exercise, and good choices"""

examples = [["mind"], ["memory"], ["sleep"],["wellness"],["nutrition"],["mnemonics"]]


def poem2_generate(word):
  p = word.lower() + "\n" + "poem using word: "
  print(f"*****Inside poem_generate - Prompt is :{p}")
  json_ = {"inputs": p,
            "parameters":
            {
            "top_p": 0.9,
          "temperature": 1.1,
          "max_new_tokens": 50,
          "return_full_text": False
          }}
  response = requests.post(API_URL, headers=headers, json=json_)
  output = response.json()
  print(f"If there was an error? Reason is : {output}")
  output_tmp = output[0]['generated_text']
  print(f"GPTJ response without splits is: {output_tmp}")
  #poem = output[0]['generated_text'].split("\n\n")[0] # +"."
  if "\n\n" not in output_tmp:
    if output_tmp.find('.') != -1:
      idx = output_tmp.find('.')
      poem = output_tmp[:idx+1]
    else:
      idx = output_tmp.rfind('\n')
      poem = output_tmp[:idx]
  else:
    poem = output_tmp.split("\n\n")[0] # +"."
  poem = poem.replace('?','')
  print(f"Poem being returned is: {poem}")
  return poem
  

def poem_generate(word):

  p = prompt + word.lower() + "\n" + "poem using word: "
  print(f"*****Inside poem_generate - Prompt is :{p}")
  json_ = {"inputs": p,
            "parameters":
            {
            "top_p": 0.9,
          "temperature": 1.1,
          "max_new_tokens": 50,
          "return_full_text": False
          }}
  response = requests.post(API_URL, headers=headers, json=json_)
  output = response.json()
  print(f"If there was an error? Reason is : {output}")
  output_tmp = output[0]['generated_text']
  print(f"GPTJ response without splits is: {output_tmp}")
  #poem = output[0]['generated_text'].split("\n\n")[0] # +"."
  if "\n\n" not in output_tmp:
    if output_tmp.find('.') != -1:
      idx = output_tmp.find('.')
      poem = output_tmp[:idx+1]
    else:
      idx = output_tmp.rfind('\n')
      poem = output_tmp[:idx]
  else:
    poem = output_tmp.split("\n\n")[0] # +"."
  poem = poem.replace('?','')
  print(f"Poem being returned is: {poem}")
  return poem

def poem_to_image(poem):
  print("*****Inside Poem_to_image")
  poem = " ".join(poem.split('\n'))
  poem = poem + " oil on canvas."
  steps, width, height, images, diversity = '50','256','256','1',15
  img = gr.Interface.load("spaces/multimodalart/latentdiffusion")(poem, steps, width, height, images, diversity)[0]
  return img

demo = gr.Blocks()

with demo:
  gr.Markdown("<h1><center>Few Shot Learning for Text - Word Image Search</center></h1>")
  gr.Markdown(
        "<div>This example uses prompt engineering to search for answers in EleutherAI large language model and follows the pattern of Few Shot Learning where you supply A 1) Task Description, 2) a Set of Examples, and 3) a Prompt.  Then few shot learning can show the answer given the pattern of the examples. More information on how it works is here:  https://huggingface.co/blog/few-shot-learning-gpt-neo-and-inference-api  Also the Eleuther AI was trained on texts called The Pile which is documented here on its github.  Review this to find what types of language patterns it can generate text for as answers:  https://github.com/EleutherAI/the-pile"
    )
  with gr.Row():
    input_word = gr.Textbox(lines=7, value=prompt)
    poem_txt = gr.Textbox(lines=7)
    output_image = gr.Image(type="filepath", shape=(256,256))
  
  b1 = gr.Button("Generate Text")
  b2 = gr.Button("Generate Image")

  b1.click(poem2_generate, input_word, poem_txt)
  b2.click(poem_to_image, poem_txt, output_image)
  #examples=examples

demo.launch(enable_queue=True)