import gradio as gr
from openai import OpenAI
import os

from transformers import pipeline
# from dotenv import load_dotenv, find_dotenv
import huggingface_hub



# _ = load_dotenv(find_dotenv()) # read local .env file
hf_token= os.environ['HF_TOKEN']
huggingface_hub.login(hf_token)

pipe = pipeline("token-classification", model='elshehawy/fine-tuned-roberta-base-ner', aggregation_strategy="first")


llm_model = 'gpt-3.5-turbo-0125'
# openai.api_key = os.environ['OPENAI_API_KEY']

client = OpenAI(
    api_key=os.environ.get("OPENAI_API_KEY"),
)


def get_completion(prompt, model=llm_model):
    messages = [{"role": "user", "content": prompt}]
    response = client.chat.completions.create(
        messages=messages,
        model=model,
        temperature=0,
    )
    return response.choices[0].message.content

def find_orgs(sentence, choice):
    prompt = f"""
    In context of named entity recognition (NER), find all organizations in the text delimited by triple backticks.
    
    text:
    ```
    {sentence}
    ```
    You should always start your answer with "Organizations are: "
    """
    if choice=='GPT':
        return get_completion(prompt)
    else:
        message = 'Organizations are:'
        org_list = []
        for ent in pipe(sentence):
            if ent['entity_group'] == 'ORG' and ent['word'] not in org_list:
                # message += f'\n- {ent["word"]} \t- score: {ent["score"]}'
                message += f'\n- {ent["word"]}'# \t- score: {ent["score"]}'
                org_list.append(ent['word'])
        return message


example_1 = """\
My latest exclusive for The Hill : Conservative frustration over Republican efforts to force a House vote on reauthorizing the Export - Import Bank boiled over Wednesday during a contentious GOP meeting.
"""
example_2 = """\
Markets open in 4 hrs 3 mins Asia drops as commodities languish, loonie dips on Canada vote Reuters – 15 minutes ago A man holding an umbrella walks in front of an electronic stock quotation board outside a brokerage in … By Shinichi Saoshiro TOKYO ( Reuters ) - Asian equities fell across the board on Tuesday after commodity prices languished in the wake of China growth woes and dampened risk sentiment, while the euro hovered near a 10 - day low ahead of a European Central Bank meeting that could open the door for more monetary easing.
"""
radio_btn = gr.Radio(choices=['GPT', 'iSemantics'], value='iSemantics', label='Available models', show_label=True)
textbox = gr.Textbox(label="Enter your text", placeholder="", lines=4)

iface = gr.Interface(fn=find_orgs, inputs=[textbox, radio_btn], outputs="text",  examples=[[example_1], [example_2]])
iface.launch(share=False)