CyNER2.0 Runtime Memory Optimized

app.py

@@ -1,12 +1,15 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
+from collections import defaultdict
 
 # Load model and tokenizer
 path_to_checkpoint = 'PranavaKailash/CyNER-2.0-DeBERTa-v3-base'
-tokenizer = AutoTokenizer.from_pretrained(path_to_checkpoint, use_fast=True, max_length=768)
+tokenizer = AutoTokenizer.from_pretrained(path_to_checkpoint, use_fast=True)
 model = AutoModelForTokenClassification.from_pretrained(path_to_checkpoint)
 ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)
 
+MAX_INPUT_LENGTH = 500  # Set an appropriate length limit
+
 def tag_sentence(sentence, entities_dict):
     """
     Add HTML tags to entities for visualization.
@@ -50,12 +53,10 @@ def perform_ner(text):
     Run NER pipeline and prepare results for display.
     """
     entities = ner_pipeline(text)
-    entities_dict = {}
+
+    entities_dict = defaultdict(list)
     for entity in entities:
-        entity_type = entity['entity']
-        if entity_type not in entities_dict:
-            entities_dict[entity_type] = []
-        entities_dict[entity_type].append({
+        entities_dict[entity['entity']].append({
             "entity": entity['entity'],
             "score": entity['score'],
             "index": entity['index'],
@@ -63,9 +64,9 @@ def perform_ner(text):
             "start": entity['start'],
             "end": entity['end']
         })
-    
+
     tagged_sentence = tag_sentence(text, entities_dict)
-    return entities_dict, tagged_sentence
+    return dict(entities_dict), tagged_sentence
 
 # Streamlit UI
 st.title("CyNER 2.0 - Named Entity Recognition")
@@ -74,7 +75,9 @@ st.write("Enter text to get named entity recognition results.")
 input_text = st.text_area("Input Text", "Type your text here...")
 
 if st.button("Analyze"):
-    if input_text.strip():
+    if len(input_text) > MAX_INPUT_LENGTH:
+        st.warning(f"Text is too long! Please enter less than {MAX_INPUT_LENGTH} characters.")
+    elif input_text.strip():
         entities_dict, tagged_sentence = perform_ner(input_text)
         
         # Display results
@@ -84,4 +87,4 @@ if st.button("Analyze"):
         st.subheader("Entities and Details")
         st.json(entities_dict)
     else:
-        st.warning("Please enter some text for analysis.")
+        st.warning("Please enter some text for analysis.")