CyNER2.0 Runtime Memory Optimized v2

app.py

@@ -4,11 +4,14 @@ 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)
+tokenizer = AutoTokenizer.from_pretrained(path_to_checkpoint, use_fast=True, max_length=768)
 model = AutoModelForTokenClassification.from_pretrained(path_to_checkpoint)
-ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)
 
-MAX_INPUT_LENGTH = 500  # Set an appropriate length limit
+# Ensure the model is loaded on CPU explicitly to avoid any device issues
+model.to('cpu')
+
+# Initialize the NER pipeline
+ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)
 
 def tag_sentence(sentence, entities_dict):
     """
@@ -53,8 +56,8 @@ def perform_ner(text):
     Run NER pipeline and prepare results for display.
     """
     entities = ner_pipeline(text)
-
     entities_dict = defaultdict(list)
+
     for entity in entities:
         entities_dict[entity['entity']].append({
             "entity": entity['entity'],
@@ -64,7 +67,7 @@ def perform_ner(text):
             "start": entity['start'],
             "end": entity['end']
         })
-
+    
     tagged_sentence = tag_sentence(text, entities_dict)
     return dict(entities_dict), tagged_sentence
 
@@ -75,9 +78,7 @@ 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 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():
+    if input_text.strip():
         entities_dict, tagged_sentence = perform_ner(input_text)
         
         # Display results
@@ -87,4 +88,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.")