app.py

import uvicorn 
import threading
from typing import Optional
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForTokenClassification
import pandas as pd
#import datasets
from pprint import pprint 

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

# Define the FastAPI app
app = FastAPI()
model_cache: Optional[object] = None

def load_model():
    
    tokenizer = AutoTokenizer.from_pretrained("LampOfSocrates/bert-cased-plodcw-sourav")
    model = AutoModelForTokenClassification.from_pretrained("LampOfSocrates/bert-cased-plodcw-sourav")
    # Mapping labels
    id2label = model.config.id2label
    # Print the label mapping
    print(f"Can recognise the following labels {id2label}")

    # Load the NER model and tokenizer from Hugging Face
    #ner_pipeline = pipeline("ner", model="dbmdz/bert-large-cased-finetuned-conll03-english")
    model = pipeline("ner", model=model, tokenizer = tokenizer)
    return model 

def load_plod_cw_dataset():
    from datasets import load_dataset
    dataset = load_dataset("surrey-nlp/PLOD-CW")
    return dataset

def get_cached_model():
    global model_cache
    if model_cache is None:
        model_cache = load_model()
    return model_cache

# Cache the model when the server starts
model = get_cached_model()



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

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

class NERRequest(BaseModel):
    text: str

@app.get("/hello")
def read_root():
    return {"message": "Hello, World!"}


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

def get_color_for_label(label: str) -> str:
    # Define a mapping of labels to colors
    color_mapping = {
        "I-LF": "red",
        "B-AC": "blue",
        "LOC": "green",
        # Add more labels and colors as needed
    }
    return color_mapping.get(label, "black")  # Default to black if label not found


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

    # Color code the entities
    color_coded_text = text
    for entity in entities:
        #print(entity)
        start, end, label = entity["start"], entity["end"], entity["entity"]
        color = get_color_for_label(label)  # You need to define this function
        entity_text = text[start:end]
        colored_entity = f'<span style="color: {color}; font-weight: bold;">{entity_text}</span>'
        color_coded_text = color_coded_text[:start] + colored_entity + color_coded_text[end:]
    
    return color_coded_text

PROJECT_INTRO = "This is a HF Spaces hosted Gradio App built by NLP Group 27 . The model has been trained on surrey-nlp/PLOD-CW dataset"

def echo(text, request: gr.Request):
    if request:
        print("Request headers dictionary:", request.headers)
        print("IP address:", request.client.host)
        print("Query parameters:", dict(request.query_params))
    return text

# Create the Gradio interface
demo = gr.Interface(
    fn=ner_demo,
    inputs=gr.Textbox(lines=10, placeholder="Enter text here..."),
    outputs="html",
    #outputs=gr.JSON(),
    title="Named Entity Recognition on PLOD-CW ",
    description=f"{PROJECT_INTRO}\n\nEnter text to extract named entities using a NER model."
)
'''
with gr.Blocks() as demo:
    gr.Markdown("# Page Title")
    gr.Markdown("## Subtitle with h2 Font")
    inputs=gr.Textbox(lines=10, placeholder="Enter text here...", label="Input Text")

    with gr.Column():
        echo_output = gr.Textbox(label="Echo Output")
        html_output = ner_demo

    with gr.Column():
        button1 = gr.Button("Submit")
'''            

CUSTOM_PATH = "/gradio"
app = gr.mount_gradio_app(app, demo, path=CUSTOM_PATH)

# Function to run FastAPI
def run_fastapi():
    uvicorn.run(app, host="0.0.0.0", port=8000)

# Function to run Gradio
def run_gradio():
    demo.launch(server_name="0.0.0.0", server_port=7860)


run_fastapi()
#threading.Thread(target=run_gradio).start()