import os

os.environ['debug'] = 'true'

import gradio as gr

from GPTagger import *
from langchain.prompts import PromptTemplate

default_prompt = """
Please understand the instructions above and do extraction in the text below.

TEXT:
\"\"\"
{text}
\"\"\"
"""

def ner(
    model: str,
    nr_calls: int,
    tag_name: str,
    tag_max_len: int,
    text: str,
    prompt: str,
    key: str,
):
    os.environ['OPENAI_API_KEY'] = key
    ner_pipeline = NerPipeline(
        tag_name=tag_name,
        nr_calls=nr_calls,
        model=model,
        tag_max_len=tag_max_len
    )
    template = PromptTemplate.from_template(prompt)

    extractions = ner_pipeline(text, template, "")

    if not extractions:
        output = []
    else:
        output = [
            {"entity": tag_name.upper(), "start": item.start, "end": item.end}
            for item in extractions
        ]
    return {"text": text, "entities": output}


with gr.Blocks() as demo:
    gr.Markdown(
        """
        # GPTagger 🏷️

        [GPTagger](https://github.com/hnliu-git/GPTagger) is a powerful text tagger that makes use of the GPT model. This tool allows you to extract tags from a given text by leveraging the capabilities of GPT.
        Simply specify the tag you want to extract from the text using prompt, you will get them highlighted in the output.
        """
    )
    with gr.Row():
        key = gr.Textbox(label='OpenAI API Key: （We don \'t record your key.)')
    with gr.Row():
        tag_name = gr.Textbox(label="Tag Name:", placeholder='Enter the tag you want to extract')
        tag_max_len = gr.Slider(
            minimum=10, maximum=1000, step=10, label="Max length of a tag", value=50
        )
    with gr.Row():
        model = gr.Dropdown(
            ["gpt-3.5-turbo-0613", "gpt-4-0613"],
            label="Model Name:",
            value="gpt-3.5-turbo-0613",
        )
        nr_call = gr.Number(label="nr_of_calls", minimum=1, value=1, precision=0)
    with gr.Row():
        prompt = gr.TextArea(
            placeholder="Enter your prompt here...",
            label="Prompt: (Please include the default prompt at the end)",
            value=default_prompt,
        )
        text = gr.TextArea(placeholder="Enter your text here...", label="Text")
    btn = gr.Button("Submit")
    output = gr.HighlightedText()
    btn.click(
        ner,
        inputs=[
            model,
            nr_call,
            tag_name,
            tag_max_len,
            text,
            prompt,
            key
        ],
        outputs=output,
    )

demo.launch()