from typing import Dict, Union import spaces from span_marker import SpanMarkerModel import gradio as gr import torch # Load the SpanMarker model from the Hub & try to put it on CUDA model = SpanMarkerModel.from_pretrained("tomaarsen/span-marker-bert-base-fewnerd-fine-super").try_cuda() # The labels, only necessary for the space documentation labels = [ "art-broadcastprogram", "art-film", "art-music", "art-other", "art-painting", "art-writtenart", "building-airport", "building-hospital", "building-hotel", "building-library", "building-other", "building-restaurant", "building-sportsfacility", "building-theater", "event-attack/battle/war/militaryconflict", "event-disaster", "event-election", "event-other", "event-protest", "event-sportsevent", "location-GPE", "location-bodiesofwater", "location-island", "location-mountain", "location-other", "location-park", "location-road/railway/highway/transit", "organization-company", "organization-education", "organization-government/governmentagency", "organization-media/newspaper", "organization-other", "organization-politicalparty", "organization-religion", "organization-showorganization", "organization-sportsleague", "organization-sportsteam", "other-astronomything", "other-award", "other-biologything", "other-chemicalthing", "other-currency", "other-disease", "other-educationaldegree", "other-god", "other-language", "other-law", "other-livingthing", "other-medical", "person-actor", "person-artist/author", "person-athlete", "person-director", "person-other", "person-politician", "person-scholar", "person-soldier", "product-airplane", "product-car", "product-food", "product-game", "product-other", "product-ship", "product-software", "product-train", "product-weapon", ] examples = [ "Amelia Earhart flew her single engine Lockheed Vega 5B across the Atlantic to Paris.", "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.", "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.", "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.", "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)", "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 .", "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).", "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.", "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.", "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.", ] @spaces.GPU def ner(text) -> Dict[str, Union[str, int, float]]: print(model.device) return { "text": text, "entities": [ { "entity": entity["label"], "word": entity["span"], "start": entity["char_start_index"], "end": entity["char_end_index"], "score": entity["score"], } for entity in model.predict(text) ], } with gr.Blocks(theme="JohnSmith9982/small_and_pretty", title="SpanMarker for Named Entity Recognition") as demo: gr.Markdown( """ # SpanMarker for Named Entity Recognition This Space showcases the [`tomaarsen/span-marker-bert-base-fewnerd-fine-super`](https://huggingface.co/tomaarsen/span-marker-bert-base-fewnerd-fine-super) [SpanMarker](https://github.com/tomaarsen/SpanMarkerNER) model for NER, which uses the [`bert-base-cased`](https://huggingface.co/bert-base-cased) encoder under the hood. The model has been trained on the 66-class supervised [FewNERD dataset](https://huggingface.co/datasets/DFKI-SLT/few-nerd) for approximately 2 hours using a slight variation of the training script from the [SpanMarker GitHub repository](https://github.com/tomaarsen/SpanMarkerNER). It reaches a very competitive 0.705 Test F1 in the all-time [FewNERD leaderboard for bert-base](https://paperswithcode.com/sota/named-entity-recognition-on-few-nerd-sup?tag_filter=155). """ ) with gr.Accordion("How to run this model locally", open=False): gr.Code( """ # pip install span_marker from span_marker import SpanMarkerModel # Load the SpanMarker model from the 🤗 Hub model = SpanMarkerModel.from_pretrained("tomaarsen/span-marker-bert-base-fewnerd-fine-super") # Optionally, place it on a GPU model.cuda() # Perform inference! 🎉 model.predict("Amelia Earhart flew her single engine Lockheed Vega 5B across the Atlantic to Paris.") # [{'span': 'Amelia Earhart', 'label': 'person-other', 'score': 0.7629689574241638, 'char_start_index': 0, 'char_end_index': 14}, # {'span': 'Lockheed Vega 5B', 'label': 'product-airplane', 'score': 0.9833564758300781, 'char_start_index': 38, 'char_end_index': 54}, # {'span': 'Atlantic', 'label': 'location-bodiesofwater', 'score': 0.7621214389801025, 'char_start_index': 66, 'char_end_index': 74}, # {'span': 'Paris', 'label': 'location-GPE', 'score': 0.9807717204093933, 'char_start_index': 78, 'char_end_index': 83}] """, language="python", ) with gr.Accordion("Full label list", open=False): gr.Markdown( "These are all 66 classes on which the model is trained. Can you come up with inputs to find all of them?\n\n" + "\n".join(f"{idx}. `{label}`" for idx, label in enumerate(labels, start=1)) ) inputs = gr.Textbox(value=examples[0], label="Text input", placeholder="Enter your text here") output = gr.HighlightedText(label="Predicted Entities") submit_btn = gr.Button("Submit") examples = gr.Examples( examples, fn=ner, inputs=inputs, outputs=output, cache_examples=True ) # Submitting inputs.submit(fn=ner, inputs=inputs, outputs=output) submit_btn.click(fn=ner, inputs=inputs, outputs=output) demo.queue() demo.launch()