# Import the required libraries
import os  # Provides a way of using operating system-dependent functionality
import requests  # For making HTTP requests to the API
import json  # For handling JSON data
from dotenv import load_dotenv, find_dotenv
import gradio as gr

# Load environment variables from .env file
load_dotenv(find_dotenv())
hf_api_key = os.getenv('HF_API_KEY')  # Hugging Face API key
API_URL = os.getenv('HF_API_NER_BASE')  # Endpoint for the NER model

# Define the `get_completion` function to interact with the Hugging Face API
def get_completion(inputs, parameters=None, endpoint_url=None):
    headers = {
        "Authorization": f"Bearer {hf_api_key}",
        "Content-Type": "application/json"
    }
    data = {"inputs": inputs}
    if parameters:
        data.update({"parameters": parameters})
    try:
        response = requests.post(endpoint_url, headers=headers, data=json.dumps(data))
        response.raise_for_status()
        return response.json()  # Return the API's JSON response
    except requests.exceptions.RequestException as e:
        print(f"Error: {e}")
        return [{"entity": "Error", "word": "Error", "score": 0}]

# Function to perform Named Entity Recognition (NER)
def ner(input):
    output = get_completion(input, parameters=None, endpoint_url=API_URL)
    return {"text": input, "entities": output}

# Create a Gradio interface
iface = gr.Interface(
    fn=ner,
    inputs=[gr.Textbox(label="Text to find entities", lines=2)],
    outputs=[gr.HighlightedText(label="Text with entities")],
    title="NER with dslim/bert-base-NER",
    description="Find entities using the `dslim/bert-base-NER` model under the hood!",
    allow_flagging="never",
    examples=["My name is Michela and I live in Italy", "My name is Andrew and work at HuggingFace"]
)

# Launch the app (this will allow you to test locally before uploading to Hugging Face)
iface.launch()