ChatGPT updates

app.py

@@ -1,14 +1,26 @@
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassification
 import gradio as gr
 
-# Load the NER pipeline
-ner_model = pipeline("ner", model="dslim/bert-base-NER", aggregation_strategy="simple")
+# Load the tokenizer and model
+try:
+    tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER")
+    model = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER")
+    ner_model = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
+except Exception as e:
+    ner_model = None
+    print(f"Error loading model: {e}")
 
 def extract_named_entities(text):
-    # Use the NER model to extract named entities
-    entities = ner_model(text)
-    # Format the output for better readability
-    return [{"Entity": ent["entity_group"], "Text": ent["word"], "Score": round(ent["score"], 3)} for ent in entities]
+    if not text.strip():
+        return [{"Entity": "N/A", "Text": "No input provided", "Score": 0.0}]
+
+    try:
+        if ner_model is None:
+            raise ValueError("Model failed to load. Check logs for details.")
+        entities = ner_model(text)
+        return [{"Entity": ent["entity_group"], "Text": ent["word"], "Score": round(ent["score"], 3)} for ent in entities]
+    except Exception as e:
+        return [{"Entity": "Error", "Text": str(e), "Score": 0.0}]
 
 # Define the Gradio interface
 iface = gr.Interface(
@@ -17,8 +29,12 @@ iface = gr.Interface(
     outputs=gr.Dataframe(headers=["Entity", "Text", "Score"], label="Named Entities"),
     title="Named Entity Recognition",
     description="Input some text and get the named entities (like names, locations, organizations).",
+    examples=[
+        ["My name is Emilio."],
+        ["Barack Obama was the President of the United States."],
+        ["OpenAI created ChatGPT."],
+    ]
 )
 
-# Launch the app
 if __name__ == "__main__":
     iface.launch()