# Copyright 2023 by Jan Philip Wahle, https://jpwahle.com/
# All rights reserved.

import os
import random
import time

import gradio as gr
import openai

openai.api_key = os.environ.get("OPENAI_API_KEY")

def create_prompt(sentence, paraphrase_type):
    prompt = {
        "messages": [
            {
                "role": "user",
                "content": f"Given the following sentence, generate a paraphrase with the following types. Sentence: {sentence}. Paraphrase Types: {paraphrase_type}"
            }
        ]
    }
    return prompt

with gr.Blocks() as demo:
    description = gr.Markdown("""
        ## Paraphrase Type Generator
        This demo uses a fine-tuned ChatGPT-3.5 model to generate paraphrases given specific paraphrase types.
        
        **How to use:**
        1. Select one or many type of paraphrase from the dropdown menu.
        2. Enter a sentence in the text box.
        3. Click the "Submit" button or hit enter.
        4. The application will generate a paraphrase of the input sentence based on the selected type.
    """)
    chatbot = gr.Chatbot()
    types = gr.Dropdown(
        ["Morphology-based changes", "Lexico-syntactic based changes", "Lexicon-based changes", "Syntax-based changes", "Discourse-based changes"],
        value="Syntax-based changes",
        multiselect=True,
        allow_custom_value=True,
    )
    msg = gr.Textbox()
    submit = gr.Button("Submit")
    clear = gr.Button("Clear")


    def user(user_message, history):
        return "", history + [[user_message, None]]

    def generate_paraphrase(user_message, paraphrase_type, history):
        history[-1][1] = ""
        prompt = create_prompt(history[-1][0], paraphrase_type)        
        bot_message = openai.ChatCompletion.create(
            model="ft:gpt-3.5-turbo-0613:personal::7xbU0xQ2",
            messages=prompt["messages"],
        )
        history[-1][1] = ""
        for character in bot_message.choices[0].message.content:
            history[-1][1] += character
            time.sleep(0.01)
            yield history

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        generate_paraphrase, [msg, types, chatbot], chatbot
    )
    submit.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        generate_paraphrase, [msg, types, chatbot], chatbot
    )
    clear.click(lambda: None, None, chatbot, queue=False)
    
demo.queue()
demo.launch()