Fixed No data in Tensor error v1

app.py

@@ -1,17 +1,28 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
 from collections import defaultdict
+import torch
 
 # 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)
-model = AutoModelForTokenClassification.from_pretrained(path_to_checkpoint)
 
-# Ensure the model is loaded on CPU explicitly to avoid any device issues
-model.to('cpu')
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(path_to_checkpoint, use_fast=True)
+
+# Load model with proper device handling
+model = AutoModelForTokenClassification.from_pretrained(
+    path_to_checkpoint,
+    torch_dtype='auto',
+    device_map='cpu'
+)
 
 # Initialize the NER pipeline
-ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)
+ner_pipeline = pipeline(
+    "ner", 
+    model=model, 
+    tokenizer=tokenizer,
+    device=-1  # Explicitly use CPU
+)
 
 def tag_sentence(sentence, entities_dict):
     """
@@ -21,10 +32,10 @@ def tag_sentence(sentence, entities_dict):
         [(e['start'], e['end'], e['entity'], e['word']) for ents in entities_dict.values() for e in ents],
         key=lambda x: x[0]
     )
-
+    
     merged_entities = []
     current_entity = None
-
+    
     for start, end, entity_type, word in all_entities:
         if current_entity is None:
             current_entity = [start, end, entity_type, word]
@@ -35,57 +46,114 @@ def tag_sentence(sentence, entities_dict):
             else:
                 merged_entities.append(tuple(current_entity))
                 current_entity = [start, end, entity_type, word]
-
+    
     if current_entity:
         merged_entities.append(tuple(current_entity))
-
+    
     tagged_sentence = ""
     last_idx = 0
-
+    
     for start, end, entity_type, _ in merged_entities:
         tagged_sentence += sentence[last_idx:start]
         entity_tag = entity_type.replace('I-', 'B-')
-        tagged_sentence += f"<span style='color:blue'><{entity_tag}></span>{sentence[start:end]}<span style='color:blue'>/{entity_tag}></span>"
+        tagged_sentence += f"<span style='color:blue; background-color: #e6f3ff; padding: 2px; border-radius: 3px;'><strong>{entity_tag}</strong></span><span style='background-color: #fff3cd; padding: 2px; border-radius: 3px;'>{sentence[start:end]}</span>"
         last_idx = end
-
+    
     tagged_sentence += sentence[last_idx:]
     return tagged_sentence
 
+@st.cache_data
 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'],
-            "score": entity['score'],
-            "index": entity['index'],
-            "word": entity['word'],
-            "start": entity['start'],
-            "end": entity['end']
-        })
-    
-    tagged_sentence = tag_sentence(text, entities_dict)
-    return dict(entities_dict), tagged_sentence
+    try:
+        entities = ner_pipeline(text)
+        entities_dict = defaultdict(list)
+        
+        for entity in entities:
+            entities_dict[entity['entity']].append({
+                "entity": entity['entity'],
+                "score": round(entity['score'], 4),
+                "index": entity['index'],
+                "word": entity['word'],
+                "start": entity['start'],
+                "end": entity['end']
+            })
+        
+        tagged_sentence = tag_sentence(text, entities_dict)
+        return dict(entities_dict), tagged_sentence
+    except Exception as e:
+        st.error(f"Error during NER processing: {str(e)}")
+        return {}, text
 
 # Streamlit UI
-st.title("CyNER 2.0 - Named Entity Recognition")
-st.write("Enter text to get named entity recognition results.")
+st.set_page_config(
+    page_title="CyNER 2.0",
+    page_icon="🔐",
+    layout="wide"
+)
 
-input_text = st.text_area("Input Text", "Type your text here...")
+st.title("🔐 CyNER 2.0 - Cybersecurity Named Entity Recognition")
+st.markdown("**Advanced NER for Cybersecurity Text Analysis**")
+st.write("Enter cybersecurity-related text to identify and extract named entities using the CyNER 2.0 model.")
 
-if st.button("Analyze"):
-    if input_text.strip():
-        entities_dict, tagged_sentence = perform_ner(input_text)
-        
-        # Display results / visualization 
-        st.subheader("Tagged Entities")
-        st.markdown(tagged_sentence, unsafe_allow_html=True)
+# Example texts
+examples = {
+    "Malware Analysis": "The Zeus trojan was detected on the victim's Windows 10 system at IP address 192.168.1.100. The malware communicated with command and control server evil.example.com using port 8080.",
+    "Vulnerability Report": "CVE-2021-44228 affects Apache Log4j versions 2.0 to 2.15.0. The vulnerability allows remote code execution through LDAP injection.",
+    "Incident Response": "Suspicious network traffic detected from IP 203.0.113.1 attempting to access /admin/login.php on our web server nginx running on Ubuntu 20.04."
+}
+
+# Sidebar for examples
+with st.sidebar:
+    st.header("Example Texts")
+    for title, text in examples.items():
+        if st.button(f"Load: {title}"):
+            st.session_state.input_text = text
+
+# Main input
+input_text = st.text_area(
+    "Input Text", 
+    value=st.session_state.get('input_text', "Enter your cybersecurity text here..."),
+    height=150,
+    key='input_text'
+)
+
+col1, col2 = st.columns([1, 4])
+with col1:
+    analyze_button = st.button("🔍 Analyze Text", type="primary")
+
+if analyze_button:
+    if input_text.strip() and input_text != "Enter your cybersecurity text here...":
+        with st.spinner("Processing text with CyNER 2.0..."):
+            entities_dict, tagged_sentence = perform_ner(input_text)
         
-        st.subheader("Entities and Details")
-        st.json(entities_dict)
+        if entities_dict:
+            # Display results
+            st.subheader("📊 Analysis Results")
+            
+            # Tagged visualization
+            st.markdown("**Tagged Entities:**")
+            st.markdown(tagged_sentence, unsafe_allow_html=True)
+            
+            # Entity summary
+            st.markdown("**Entity Summary:**")
+            entity_counts = {k: len(v) for k, v in entities_dict.items()}
+            
+            cols = st.columns(min(len(entity_counts), 4))
+            for i, (entity_type, count) in enumerate(entity_counts.items()):
+                with cols[i % 4]:
+                    st.metric(entity_type.replace('B-', '').replace('I-', ''), count)
+            
+            # Detailed results
+            with st.expander("📋 Detailed Entity Information", expanded=False):
+                st.json(entities_dict)
+        else:
+            st.info("No entities detected in the provided text.")
     else:
-        st.warning("Please enter some text for analysis.")
+        st.warning("⚠️ Please enter some text for analysis.")
+
+# Footer
+st.markdown("---")
+st.markdown("**CyNER 2.0** - Powered by DeBERTa-v3-base | Built with Streamlit")

requirements.txt

@@ -1,2 +1,5 @@
-transformers
-torch
+streamlit>=1.28.0
+transformers>=4.30.0
+torch>=2.0.0
+tokenizers>=0.13.0
+numpy>=1.21.0