import gradio as gr
from transformers import pipeline
from fastapi import FastAPI
from pydantic import BaseModel
from typing import List, Dict

# Initialize the NER pipeline
ner_model = pipeline("ner", grouped_entities=True)

# Define the FastAPI app
app = FastAPI()

# Define the request and response models for the API
class NERRequest(BaseModel):
    text: str

class Entity(BaseModel):
    entity_group: str
    start: int
    end: int
    score: float
    word: str

class NERResponse(BaseModel):
    entities: List[Entity]

@app.post("/ner", response_model=NERResponse)
def get_entities(request: NERRequest):
    # Use the NER model to detect entities
    entities = ner_model(request.text)
    # Convert entities to the response model
    response_entities = [Entity(**entity) for entity in entities]
    return NERResponse(entities=response_entities)

# Define the Gradio interface function
def ner_demo(text):
    entities = ner_model(text)
    return {"entities": entities}

# Create the Gradio interface
iface = gr.Interface(
    fn=ner_demo,
    inputs=gr.Textbox(lines=10, placeholder="Enter text here..."),
    outputs=gr.JSON(),
    title="Named Entity Recognition",
    description="Enter text to extract named entities using a NER model."
)

# Launch the Gradio interface
iface.launch(share=True)