import gradio as gr
from SynonymEditor import SynonymEditor

LANGS = ["de", "en"]


def replace_synonyms(api_key, language, text, use_few_shots, few_shots, temperature):
    model_engine = "text-davinci-003"
    max_tokens = 500
    editor = SynonymEditor(api_key, model_engine,
                           max_tokens, temperature, language)
    if (use_few_shots):
        return editor._edit_text(text, few_shots)
    else:
        return editor._edit_text(text)


api_key = gr.inputs.Textbox(label="OpenAI API Key", lines=1,
                            default="YOUR_API_KEY_HERE")

language = gr.inputs.Dropdown(
    label="Language", choices=LANGS, default="de")

use_few_shots = gr.inputs.Checkbox(label="Use Few Shots", default=False)

few_shots = gr.inputs.Textbox(
    label="Few Shots", lines=10, default='Input: The cat sat on the mat. Output: The cat sat on the rug. \nInput: "She walked to the store to buy some milk.", she said to her mother. Output: "She walked to the shop to buy some milk.", she said to her mother.\nInput: He cooked a "perfect" dinner for his family. Output: He prepared a "perfect" dinner for his family.\nInput: The sun was shining brightly in the sky. Output: The sun was beaming brightly in the sky.\nInput: Little Red Riding Hood was enjoying the warm summer day so much, that she didn\'t notice a sinister shadow approaching out of the forest behind her...	 Output: Little Red Riding Hood was enjoying the warm summer day so much, that she didn\'t notice a dark shadow approaching out of the forest behind her...')

input_text = gr.inputs.Textbox(
    label="Input Text", lines=10, default="Enter your text here.")

temperature_slider = gr.inputs.Slider(
    label="Temperature", minimum=0, maximum=1, default=0.7, step=0.1)

output_text = gr.outputs.Textbox(label="Output Text")

io = gr.Interface(fn=replace_synonyms, inputs=[api_key, language, input_text, use_few_shots, few_shots, temperature_slider], outputs=output_text, title="Synonym Replacer",
                  description="Replace words in a text with their synonyms.")
io.launch()