app.py

import json
from typing import Dict, Union

import gradio as gr
import requests
import uvicorn
from fastapi import FastAPI
from gliner import GLiNER
from starlette.responses import StreamingResponse, JSONResponse


model = GLiNER.from_pretrained("urchade/gliner_medium-v2.1")
model.eval()

examples = [
    [
        """
* Data Scientist, Data Analyst, or Data Engineer with 1+ years of experience.
* Experience with technologies such as Docker, Kubernetes, or Kubeflow
* Machine Learning experience preferred
* Experience with programming languages such as Python, C++, or SQL preferred
* Experience with technologies such as Databricks, Qlik, TensorFlow, PyTorch, Python, Dash, Pandas, or NumPy preferred
* BA or BS degree
* Active Secret OR Active Top Secret or Active TS/SCI clearance
""",
        "software package, programing language, software tool, degree, job title",
        0.3,
        False,
    ],
    [
        "Libretto by Marius Petipa, based on the 1822 novella ``Trilby, ou Le Lutin d'Argail`` by Charles Nodier, first presented by the Ballet of the Moscow Imperial Bolshoi Theatre on January 25/February 6 (Julian/Gregorian calendar dates), 1870, in Moscow with Polina Karpakova as Trilby and Ludiia Geiten as Miranda and restaged by Petipa for the Imperial Ballet at the Imperial Bolshoi Kamenny Theatre on January 17–29, 1871 in St. Petersburg with Adèle Grantzow as Trilby and Lev Ivanov as Count Leopold.",
        "person, book, location, date, actor, character",
        0.3,
        True,
    ],
    [
        "However, both models lack other frequent DM symptoms including the fibre-type dependent atrophy, myotonia, cataract and male-infertility.",
        "disease, symptom",
        0.3,
        False,
    ],
    [
        "Synergy between signal transduction pathways is obligatory for expression of c-fos in B and T cell lines: implication for c-fos control via surface immunoglobulin and T cell antigen receptors.",
        "DNA, RNA, cell line, cell type, protein",
        0.3,
        False,
    ],
    [
        "The choice of the encoder and decoder modules of dnpg can be quite flexible, for instance long short term memory networks (lstm) or convolutional neural network (cnn).",
        "short acronym, long acronym",
        0.3,
        False,
    ],
    [
        "Amelia Earhart flew her single engine Lockheed Vega 5B across the Atlantic to Paris.",
        "person, company, location, airplane",
        0.3,
        True,
    ],
    [
        "Feldman is a contributor to NBC Sports Boston's ``State of the Revs`` and ``Revolution Postgame Live`` programs as well as to 98.5 the SportsHub, SiriusXM FC's MLS coverage and to other New England and national radio outlets and podcasts.",
        "person, company, location",
        0.3,
        False,
    ],
    [
        "On 25 July 1948, on the 39th anniversary of Bleriot's crossing of the English Channel, the Type 618 Nene-Viking flew Heathrow to Paris (Villacoublay) in the morning carrying letters to Bleriot's widow and son (secretary of the FAI), who met it at the airport.",
        "date, location, person, organization",
        0.3,
        False,
    ],
    [
        "Leo & Ian won the 1962 Bathurst Six Hour Classic at Mount Panorama driving a Daimler SP250 sports car, (that year the 500 mile race for touring cars were held at Phillip Island)",
        "person, date, location, organization, competition",
        0.3,
        False,
    ],
    [
        "The Shore Line route of the CNS & M until 1955 served, from south to north, the Illinois communities of Chicago, Evanston, Wilmette, Kenilworth, Winnetka, Glencoe, Highland Park, Highwood, Fort Sheridan, Lake Forest, Lake Bluff, North Chicago, Waukegan, Zion, and Winthrop Harbor as well as Kenosha, Racine, and Milwaukee (the ``KRM'') in Wisconsin.",
        "location, organization, date",
        0.3,
        False,
    ],
    [
        "Comet C/2006 M4 (SWAN) is a non-periodic comet discovered in late June 2006 by Robert D. Matson of Irvine, California and Michael Mattiazzo of Adelaide, South Australia in publicly available images of the Solar and Heliospheric Observatory (SOHO).",
        "person, organization, date, location",
        0.3,
        False,
    ],
    [
        "From November 29, 2011 to March 31, 2012, Karimloo returned to ``Les Misérables`` to play the lead role of Jean Valjean at The Queen's Theatre, London, for which he won the 2013 Theatregoers' Choice Award for Best Takeover in a Role.",
        "person, actor, award, date, location",
        0.3,
        False,
    ],
    [
        "A Mexicali health clinic supported by former Baja California gubernatorial candidate Enrique Acosta Fregoso (PRI) was closed on June 15 after selling a supposed COVID-19 ``cure'' for between MXN $10,000 and $50,000.",
        "location, organization, person, date, currency",
        0.3,
        False,
    ],
    [
        "Built in 1793, it was the home of Mary Young Pickersgill when she moved to Baltimore in 1806 and the location where she later sewed the ``Star Spangled Banner'', in 1813, the huge out-sized garrison flag that flew over Fort McHenry at Whetstone Point in Baltimore Harbor in the summer of 1814 during the British Royal Navy attack in the Battle of Baltimore during the War of 1812.",
        "date, person, location, organization, event, flag",
        0.3,
        False,
    ],
]


def ner(text: str, labels: str, threshold: float, nested_ner: bool) -> Dict[str, Union[str, int, float]]:
    return [
        {
            "entity": entity["label"],
            "word": entity["text"],
            "start": entity["start"],
            "end": entity["end"],
        }
        for entity in model.predict_entities(
            text, [label.strip() for label in labels.split(",")], flat_ner=not nested_ner, threshold=threshold
        )
    ]


demo = gr.Interface(
    fn=ner,
    inputs=[
        gr.Textbox(label="Text input", placeholder="Enter your text here"),
        gr.Textbox(
            label="Labels",
            placeholder="Enter your labels here (comma separated)",
        ),
        gr.Slider(
            0,
            1,
            value=0.3,
            step=0.01,
            label="Threshold",
            info="Lower the threshold to increase how many entities get predicted.",
        ),
        gr.Checkbox(
            value=False,
            label="Nested NER",
            info="Allow for nested NER?",
        )
    ],
    outputs=[gr.JSON()],
    examples=examples,
    cache_examples="lazy",
    title="NER.tool",
    description="Named Entity Recognition using urchade/gliner_medium-v2.1",
    allow_flagging="never",
)

app = FastAPI()

@app.head("/ner")
def predict_head():
    return StreamingResponse("", media_type="application/json")

@app.get("/ner")
def predict_get(text: str = "", labels: str = "", threshold: float = 0.3, nested_ner: bool = False):
    predict_response = requests.post('http://localhost:7860/call/predict', json={'data': [text, labels, threshold, nested_ner]}).json()
    if "event_id" not in predict_response:
        return predict_response
    return JSONResponse(json.loads(requests.get(f'http://localhost:7860/call/predict/{predict_response["event_id"]}').text.split("data: ", 1)[-1]))

if __name__ == "__main__":
    app = gr.mount_gradio_app(app, demo, path="/")
    uvicorn.run(app, host="0.0.0.0", port=7860)