import gradio as gr
import os
import openai

openai.organization = os.environ.get('ORGANIZATION')
openai.api_key = os.environ.get('API_KEY')

def classificate_with_gpt3(text, labels, engine="text-similarity-davinci-001"):
  prompt = f"""The following text may fall inside one of this categories: {labels}, None

  {text}
  
  Category:"""
  response = openai.Completion.create(
    model=engine,
    prompt=prompt,
    temperature=0.7,
    max_tokens=10,
    stop=["\n"],
    top_p=1,
    frequency_penalty=0,
    presence_penalty=0
  )
  completion = response['choices'][0]['text'].strip()
  return completion
 
def gpt3_zero_shot_classification(text, labels):  
  completion = classificate_with_gpt3(text, labels, engine="text-babbage-001")
  return completion

# Gradio UI

with gr.Blocks(css=".gradio-container { background-color: white; }") as demo:
  gr.Markdown(f"# GPT-3 Zero shot classification app")
  with gr.Row():
    context = gr.Textbox(lines=3, label="Context", placeholder="Context here...")
  with gr.Row():
    labels = gr.Textbox(lines=3, label="Labels (Comma separated)", placeholder="Labels here...")
  btn = gr.Button(value="Analyze!", variant="primary")
  with gr.Row():
    out = gr.Textbox()
  btn.click(fn=gpt3_zero_shot_classification, inputs=[context,labels], outputs=[out])
  gr.Examples(
    [
      [ "On 22 February 2014, Ukrainian president Viktor Yanukovych was ousted from office as a result of the Euromaidan and the Revolution of Dignity, which broke out after his decision to reject the European Union–Ukraine Association Agreement and instead pursue closer ties with Russia and the Eurasian Economic Union. Shortly after Yanukovych's overthrow and exile to Russia, Ukraine's eastern and southern regions erupted with pro-Russia unrest.", "Entertainment, Business, Politics" ],
      [ "This dull recreation of the animated film doesn’t strive for anything more than what was contained in the original version of this film and actually delivers less.", "Entertainment, Business, Politics" ],
      [ "The third beta of iOS 16.1 that was released earlier this week expands the Adaptive Transparency feature introduced with the second-generation AirPods Pro to the original ‌AirPods Pro‌.", "Entertainment, Business, Politics" ],
      [ "To make this into soup, bring 3 1/2 cups water to a boil in a medium saucepan. Whisk in soup mix, reduce heat to low, and simmer for 5 minutes.", "Entertainment, Business, Politics" ]
    ],
    [context, labels],
    fn=gpt3_zero_shot_classification
  )

demo.launch(debug=True)