import os
import torch
import gradio as gr
from huggingface_hub import login
from extractor import EntityExtractor

def load_entity_extractor(model_name: str):
    return EntityExtractor(model_name)

# Authenticate with Hugging Face using token from environment variable
login(token=os.getenv('HF_TOKEN'))

# Load the model
EXTRACTOR_MODEL = "fastino/buildcom-ner-v0.2"
extractor_model = load_entity_extractor(EXTRACTOR_MODEL)

def extract_entities_with_types(text, entity_types_str):
    entity_types = [etype.strip() for etype in entity_types_str.split(',') if etype.strip()]
    output_entity = extractor_model.extract_entities(text, entity_types)
    return output_entity

examples = [
    ["Find a sofa with walnut finish and modern style", "brand, color_finish, style, collection, dimension, feature, product_type, part_number"],
    ["Order a red leather chair from the luxury collection", "brand, color_finish, style, collection, dimension, feature, product_type, part_number"],
    ["Looking for a rectangular dining table in oak wood", "brand, color_finish, style, collection, dimension, feature, product_type, part_number"],
    ["Searching for an ergonomic office chair in black", "brand, color_finish, style, collection, dimension, feature, product_type, part_number"],
    ["Buy a round marble coffee table from the latest collection", "brand, color_finish, style, collection, dimension, feature, product_type, part_number"],
    ["Get a modular sofa in gray with built-in storage", "brand, color_finish, style, collection, dimension, feature, product_type, part_number"],
    ["I need a king-size bed from the premium collection", "brand, color_finish, style, collection, dimension, feature, product_type, part_number"],
    ["Order a brass chandelier from the heritage style collection", "brand, color_finish, style, collection, dimension, feature, product_type, part_number"],
    ["Find a blue velvet ottoman with a gold finish", "brand, color_finish, style, collection, dimension, feature, product_type, part_number"]
]

# Create Gradio interface
iface = gr.Interface(
    fn=extract_entities_with_types,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter text here...", label="Text"),
        gr.Textbox(
            lines=1,
            placeholder="Enter entity types separated by commas (e.g., 'brand, color_finish')", 
            label="Entity Types",
            value="brand, color_finish, style, collection, dimension, feature, product_type, part_number"
        )
    ],
    outputs="json",
    title="Entity Extraction for Product Specifications",
    description=(
        "Extract entities from the input text. Specify the types of entities you'd like to extract using a comma-separated list. "
        "For example, you can enter 'brand, product_type' to extract those entities."
    ),
    examples=examples
)

if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)