Add application file

app.py

@@ -1,7 +1,113 @@
 import gradio as gr
+import spaces
+from transformers import AutoTokenizer, AutoModelForTokenClassification
+from transformers import pipeline
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Model names (keeping it programmatic)
+model_names = [
+    "hadiaskari98/Software_NER_prod",
+    "hadiaskari98/Software_UniNER",
+    "hadiaskari98/Hardware_UniNER",
+    "hadiaskari98/Vulnerability_UniNER",
+]
+
+example_sent = (
+    "The Ingress concept lets you map traffic to different backends based on rules you define via the Kubernetes API."
+)
+
+# Programmatically build the model info dict
+model_info = {
+    model_name: {
+        "link": f"https://huggingface.co/{model_name}",
+        "usage": f"""from transformers import pipeline
+ner = pipeline("ner", model=AutoModelForTokenClassification.from_pretrained("{model_name}"), tokenizer=AutoTokenizer.from_pretrained("{model_name}"))
+result = ner("{example_sent}")
+print(result)""",
+    }
+    for model_name in model_names
+}
+
+# Load models into a dictionary programmatically for the analyze function
+models = {
+    model_names[0]: pipeline("ner", model=AutoModelForTokenClassification.from_pretrained(model_names[0]),
+                             tokenizer=AutoTokenizer.from_pretrained(model_names[0]), device=0)
+}
+
+
+# Function to display model info (link and usage code)
+def display_model_info(model_name):
+    info = model_info[model_name]
+    usage_code = info["usage"]
+    link_button = f'[Open model page for {model_name} ]({info["link"]})'
+    return usage_code, link_button
+
+
+# Function to run NER on input text
+@spaces.GPU
+def analyze_text(text, model_name):
+    if model_name not in models:
+        models.clear()
+        models[model_name] = pipeline("ner", model=AutoModelForTokenClassification.from_pretrained(model_name),
+                                      tokenizer=AutoTokenizer.from_pretrained(model_name), device=0)
+    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"]
+        # 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("# Secure Chain Named Entity Recognition (NER)")
+
+    # Dropdown for model selection
+    model_selector = gr.Dropdown(
+        choices=list(model_info.keys()),
+        value=list(model_info.keys())[0],
+        label="Select Model",
+    )
+
+    # Textbox for input text
+    text_input = gr.Textbox(
+        label="Enter Text",
+        lines=5,
+        value=example_sent,
+    )
+    analyze_button = gr.Button("Run NER Model")
+    output = gr.HighlightedText(label="NER Result", combine_adjacent=True)
+
+    # Outputs: usage code, model page link, and analyze button
+    code_output = gr.Code(label="Use this model", visible=True)
+    link_output = gr.Markdown(
+        f"[Open model page for {model_selector} ]({model_selector})"
+    )
+    # Button for analyzing the input text
+    analyze_button.click(
+        analyze_text, inputs=[text_input, model_selector], outputs=output
+    )
+
+    # Trigger the code output and model link when model is changed
+    model_selector.change(
+        display_model_info, inputs=[model_selector], outputs=[code_output, link_output]
+    )
+
+    # Call the display_model_info function on load to set initial values
+    demo.load(
+        fn=display_model_info,
+        inputs=[model_selector],
+        outputs=[code_output, link_output],
+    )
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
 demo.launch()

requirements.txt

@@ -0,0 +1,5 @@
+transformers
+gradio
+torch
+sentencepiece
+spaces