Upload folder using huggingface_hub

app.py

@@ -1,29 +1,60 @@
-from transformers import pipeline
 import gradio as gr
+from transformers import pipeline
 
-
-# List of NER models
-models = ["dslim/bert-base-NER", "dslim/bert-base-NER-uncased", "dslim/bert-large-NER"]
-
-
-def ner(text, model_choice):
-    ner_pipeline = pipeline("ner", model=model_choice)
-    output = ner_pipeline(text)
-    return {"text": text, "entities": output}
-
-
-examples = [
-    "Does Chicago have any stores and does Joe live here?",
-]
-
-demo = gr.Interface(
-    fn=ner,
-    inputs=[
-        gr.Textbox(placeholder="Enter sentence here..."),
-        gr.Dropdown(choices=models, label="Choose NER Model"),
-    ],
-    outputs=gr.HighlightedText(),
-    examples=examples,
-)
+# Load the NER models
+# Load the NER models
+models = {
+    "dslim/bert-base-NER": pipeline(
+        "ner", model="dslim/bert-base-NER", grouped_entities=True
+    ),
+    "dslim/bert-base-NER-uncased": pipeline(
+        "ner", model="dslim/bert-base-NER-uncased", grouped_entities=True
+    ),
+    "dslim/bert-large-NER": pipeline(
+        "ner", model="dslim/bert-large-NER", grouped_entities=True
+    ),
+    "dslim/distilbert-NER": pipeline(
+        "ner", model="dslim/distilbert-NER", grouped_entities=True
+    ),
+}
+
+
+def process(text, model_name):
+    ner = models[model_name]
+    ner_results = ner(text)
+    highlighted_text = []
+    last_idx = 0
+    for entity in ner_results:
+        start = entity["start"]
+        end = entity["end"]
+        label = entity["entity_group"]
+        # Add non-entity text
+        if start > last_idx:
+            highlighted_text.append((text[last_idx:start], None))
+        # Add entity text
+        highlighted_text.append((text[start:end], label))
+        last_idx = end
+    # Add any remaining text after the last entity
+    if last_idx < len(text):
+        highlighted_text.append((text[last_idx:], None))
+    return highlighted_text
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Named Entity Recognition with BERT Models")
+    with gr.Row():
+        model_selector = gr.Dropdown(
+            choices=list(models.keys()),
+            value=list(models.keys())[0],
+            label="Select Model",
+        )
+    text_input = gr.Textbox(
+        label="Enter Text",
+        lines=5,
+        value="Hugging Face Inc. is a company based in New York City. Its headquarters are in DUMBO, therefore very close to the Manhattan Bridge.",
+    )
+    output = gr.HighlightedText(label="Named Entities")
+    analyze_button = gr.Button("Analyze")
+    analyze_button.click(process, inputs=[text_input, model_selector], outputs=output)
 
 demo.launch()