Update app.py

app.py

@@ -1,103 +1,78 @@
 from typing import Dict, Union
 from gliner import GLiNER
 import gradio as gr
+import os
 
-model = GLiNER.from_pretrained("urchade/gliner_multi_pii-v1")
+model = GLiNER.from_pretrained(
+    "gravitee-io/gliner-pii-detection",
+    token=os.getenv("HUGGINGFACE_TOKEN"),
+    load_onnx_model=True,
+    load_tokenizer=True, onnx_model_file="model.onnx"
+)
 
 examples = [
     [
-        "Pierre Dubois, résident de Paris, a fondé sa propre entreprise, Le Petit Café, située au 15 Rue de la Paix. Son numéro d'entreprise est FR-987654321-1, et il utilise le compte bancaire 9876543210 pour les transactions.",
-        "person, organization, address, company registration number, bank account number",
+        "Jana Kowalczyk's driver license number is PL-DL-55443322 and she resides at 78 Ulica Nowowiejska, Wrocław. Her contact email is jana.k@example.com.",
+        "name, driver_license_number, street_address, email",
         0.5,
         False,
     ],
     [
-        "Leticia Ramírez, una habitante de Barcelona, tiene una cita médica programada en el Hospital General de Cataluña, situado en 10 Calle de los Ángeles. Su número de la seguridad social es ES-123456789-A y su grupo sanguíneo es AB+.",
-        "person, location, address, social security number, blood type",
+        "Nguyen Van Long from Hanoi logs in from the IP 10.0.0.5 and uses the API key: 12ab34cd56ef78gh90ij. His company is VietNet Global.",
+        "name, street_address, ipv4, api_key, company",
         0.5,
         False,
     ],
     [
-        "John Smith, from London, teaches mathematics at Royal Academy located at 25 King’s Road. His employee ID is UK-987654-321 and he has been working there since 2015.",
-        "person, profession, organization, address, employee ID number",
+        "Sarah Johnson made a transaction using the credit card 379354508162306 and the CVV 834. Her IBAN is GB29 NWBK 6016 1331 9268 19.",
+        "name, credit_card_number, credit_card_security_code, iban",
         0.5,
         False,
     ],
     [
-        "In Frankfurt, Claudia Weber frequently visits her local bank branch, Deutsche Bank, at 48 Hauptstraße. Her account number is DE-1234567890123456, used primarily for her mortgage payments.",
-        "person, location, address, bank account number",
+        "Employee Thomas Becker has the ID DE-EMP-44991 and joined DataFlux GmbH on 2021-12-01. His internal email is t.becker@dataflux.de.",
+        "name, employee_id, company, date, email",
         0.5,
         False,
     ],
     [
-        "Marta Rossi, residente a Roma, ha acquistato un appartamento al 123 Via Condotti. Il numero di registrazione della proprietà è IT-654321-2018 e il mutuo è gestito tramite la Banca d'Italia con numero di conto 3216549870.",
-        "person, address, property registration number, bank account number",
+        "Laura Rossi lives at Via Roma 101, Milano. Her social security number is IT-9988776655 and she was born on 1982-07-14.",
+        "name, street_address, ssn, date_of_birth",
         0.5,
         False,
     ],
     [
-        "Paulo Coelho, um turista do Brasil, fez um seguro de viagem com a empresa Seguros PT antes de sua viagem para Lisboa. O número da apólice é BR-987654321-123 e inclui cobertura médica.",
-        "person, nationality, company, insurance policy number, coverage",
+        "Omar El-Zein uses the SWIFT code BOFAUS3N and his bank routing number is 026009593. He lives near 12 Al-Azhar Street, Cairo.",
+        "name, swift_bic_code, bank_routing_number, street_address",
         0.5,
         False,
     ],
     [
-        "Julia Fischer, eine Kundin aus München, hat bei der BayWa AG, einem großen Anbieter von Baustoffen mit Sitz am 77 Industriestraße, einen Kredit aufgenommen. Die Kreditnummer lautet DE-12345678.",
-        "person, city, organization, address, loan number",
+        "Chen Wei's employee badge shows ID EMP-CN-8899. He signed the contract on 2023-03-20 at 10:15 AM using the password Dragon@123.",
+        "name, employee_id, date, time, password",
         0.5,
         False,
     ],
     [
-        "Carlos Sánchez, profesor en la Universidad de Madrid, reside en el 5 Calle de Alcalá. Su número de identificación de profesor es ES-192837465 y tiene un doctorado en filosofía.",
-        "person, profession, address, teacher ID number, degree",
+        "Fatoumata Diarra, born 1994-04-04, lives at 45 Avenue de la Liberté, Bamako. Her BBAN is ML2930012345678901234567890.",
+        "name, date_of_birth, street_address, bban",
         0.5,
         False,
     ],
     [
-        "Sophie Dupont, une journaliste française, travaille pour Le Monde, basé au 33 rue des Écoles à Paris. Son numéro d'identification de presse est FR-75649023.",
-        "person, profession, organization, address, press ID number",
+        "Daniel Evans has a passport number K01234567 and a permanent address at 500 Pine Street, Seattle. His contact number is +1-206-555-0199.",
+        "name, passport_number, street_address, phone_number",
         0.5,
         False,
     ],
     [
-        "Manuel Oliveira, um agricultor em Porto, possui uma grande plantação de vinhas na Rua da Estrada, 120. O número de registro agrícola é PT-5678912345.",
-        "person, profession, address, agricultural registration number",
+        "Alejandro Torres created his customer account on 2024-08-30 using email ale.torres@correo.mx and ID CUST-MX-1122.",
+        "name, date, email, customer_id",
         0.5,
         False,
     ],
-    [
-        "Elisa Müller, eine Künstlerin aus Berlin, hat ihre neueste Skulptur im öffentlichen Park am Alexanderplatz ausgestellt. Ihre Künstlernummer lautet DE-112233445.",
-        "person, profession, location, artist ID number",
-        0.5,
-        False,
-    ],
-    [
-        "Federico García, un jugador de fútbol de Sevilla, ha firmado un contrato de tres años con el club Real Betis. Su número de licencia deportiva es ES-9876543210.",
-        "person, profession, organization, sports license number",
-        0.5,
-        False,
-    ],
-    [
-        "Sarah White, a London-based actress, will be performing in 'Hamlet' at the Globe Theatre located at 21 New Globe Walk. Her Equity membership number is UK-1234567.",
-        "person, profession, location, address, membership number",
-        0.5,
-        False,
-    ],
-    [
-        "Ricardo Mello, engenheiro civil, trabalha na construção da nova barragem no Rio Douro, Portugal. Seu número de registro profissional é PT-987654321.",
-        "person, profession, project location, professional registration number",
-        0.5,
-        False,
-    ],
-    [
-        "Giuseppe Conti, un cliente di Milano, ha fatto un acquisto presso il negozio La Rinascente situato in Piazza Duomo. Il numero della sua carta di credito è IT-4567891234567891.",
-        "person, location, address, credit card number",
-        0.5,
-        False,
-    ]
 ]
 
-
 def ner(
     text, labels: str, threshold: float, nested_ner: bool
 ) -> Dict[str, Union[str, int, float]]:
@@ -122,19 +97,11 @@ def ner(
 with gr.Blocks(title="GLiNER-M-v2.1") as demo:
     gr.Markdown(
         """
-        # GLiNER-PII (Personnally Identifiable Information extraction)
+        # Gravitee PII (Personnally Identifiable Information extraction)
 
         GLiNER is a Named Entity Recognition (NER) model capable of identifying any entity type using a bidirectional transformer encoder (BERT-like). It provides a practical alternative to traditional NER models, which are limited to predefined entities, and Large Language Models (LLMs) that, despite their flexibility, are costly and large for resource-constrained scenarios.
-
-        The model has been trained by fine-tuning urchade/gliner_multi-v2.1 on the urchade/synthetic-pii-ner-mistral-v1 dataset.
-        
-        ## Links
-
-        * Model: https://huggingface.co/urchade/gliner_multi_pii-v1
-        * All GLiNER models: https://huggingface.co/models?library=gliner
-        * Paper: https://arxiv.org/abs/2311.08526
-        * Repository: https://github.com/urchade/GLiNER
         """
+
     )
     with gr.Accordion("How to run this model locally", open=False):
         gr.Markdown(
@@ -150,31 +117,33 @@ with gr.Blocks(title="GLiNER-M-v2.1") as demo:
             """
         )
         gr.Code(
-            '''
-from gliner import GLiNER
-
-model = GLiNER.from_pretrained("urchade/gliner_multi_pii-v1")
-
-text = """
-Harilala Rasoanaivo, un homme d'affaires local d'Antananarivo, a enregistré une nouvelle société nommée "Rasoanaivo Enterprises" au Lot II M 92 Antohomadinika. Son numéro est le +261 32 22 345 67, et son adresse électronique est harilala.rasoanaivo@telma.mg. Il a fourni son numéro de sécu 501-02-1234 pour l'enregistrement.
-"""
+            """
+model = GLiNER.from_pretrained(
+    "gravitee-io/gliner-pii-detection",
+    load_onnx_model=True,
+    load_tokenizer=True, onnx_model_file="model.onnx"
+)
+
+text = '''
+Hey, just a quick update. I talked to David yesterday. 
+He sent over the files from his private email (david.doe@example.com), and we should be careful with his SSN: 123-45-6789.
+Also, please don't push the GitHub repo until we remove the API key: ghp_abcdEfgh1234567890.
+He mentioned his new address is 123 Maple Street in New York. 
+His PC adress is 192.168.1.100.
+'''
+
+labels = ["name",
+          "email",
+          "ssn",
+          "api_key",
+          "street_address",
+          "date",
+          "ipv4"]
 
-labels = ["work", "booking number", "personally identifiable information", "driver licence", "person", "book", "full address", "company", "actor", "character", "email", "passport number", "Social Security Number", "phone number"]
 entities = model.predict_entities(text, labels)
 
 for entity in entities:
-    print(entity["text"], "=>", entity["label"])
-            ''',
-            language="python",
-        )
-        gr.Code(
-            """
-Harilala Rasoanaivo => person
-Rasoanaivo Enterprises => company
-Lot II M 92 Antohomadinika => full address
-+261 32 22 345 67 => phone number
-harilala.rasoanaivo@telma.mg => email
-501-02-1234 => Social Security Number
+    print(entity["text"], "=>", entity["label"], "=>", entity["score"])
             """
         )