import gradio as gr
from transformers import pipeline
from huggingface_hub import login
import os

token = os.getenv("HF_TOKEN")  # Get token from environment variable
if token:
    login(token=token)
else:
    print("⚠️ Warning: HF_TOKEN not found. Set it in Space settings if using private models.")


if token:
    login(token=token)

# Load your custom NER model (unchanged)
ner_pipeline = pipeline(
    "token-classification",
    model="gamalyxd/outlaw-ocean-ner",
    aggregation_strategy="simple"
)

# Optional: merge consecutive spans of the same entity (unchanged)
def merge_spans(entities):
    merged = []
    for ent in entities:
        if (merged
            and ent['entity_group'] == merged[-1]['entity_group']
            and ent['start'] == merged[-1]['end']):
            # Extend previous span
            merged[-1]['word'] += ent['word']
            merged[-1]['end'] = ent['end']
            merged[-1]['score'] = max(merged[-1]['score'], ent['score'])
        else:
            merged.append(ent.copy())
    for ent in merged:
        ent['word'] = ent['word'].strip()
    return merged

# Function to run NER (unchanged)
def run_ner(text):
    results = ner_pipeline(text)
    results = merge_spans(results)
    if not results:
        return "No entities detected."
    return "\n".join(
        f"{ent['entity_group']:12s} ({ent['score']:.2f}): \"{ent['word']}\""
        for ent in results
    )

# Gradio interface (unchanged)
demo = gr.Interface(
    fn=run_ner,
    inputs=gr.Textbox(lines=8, placeholder="Enter maritime or fisheries text here..."),
    outputs=gr.Textbox(label="Detected Entities"),
    title="🌊 Outlaw Ocean NER Model",
    description="Identifies key entities in maritime and fisheries reports, including vessels, labor issues, IUU fishing, ESG claims, and more."
)

demo.launch()