Update app.py

app.py

@@ -12,12 +12,17 @@ from kokoro import KPipeline
 import re
 import json
 import nltk
+import stanza
 from transformers import BartForConditionalGeneration, BartTokenizer
 from nltk.tokenize import sent_tokenize
 
 nltk.download("punkt")
 nltk.download("punkt_tab")
 
+# Load Stanza's NER model
+stanza.download("en")  # Load English pipeline (can be changed for other languages)
+nlp = stanza.Pipeline("en", processors="tokenize,ner", use_gpu=False)  # Disable GPU for Hugging Face Spaces
+
 # Initialize KokoroTTS with default English
 kokoro_tts = KPipeline(lang_code='a', device="cpu")  # Load initially on CPU
 
@@ -136,6 +141,36 @@ def detect_language(text):
     except:
         return "en"  # Default to English if detection fails
 
+### 2️⃣ Named Entity Recognition (NER) Using Stanza
+def extract_entities_with_stanza(text, chunk_size=1000):
+    """Splits text into chunks, runs Stanza NER, and combines results."""
+    sentences = sent_tokenize(text)
+    chunks = []
+    current_chunk = []
+    current_length = 0
+
+    # Split text into manageable chunks
+    for sentence in sentences:
+        if current_length + len(sentence) > chunk_size:
+            chunks.append(" ".join(current_chunk))
+            current_chunk = [sentence]
+            current_length = len(sentence)
+        else:
+            current_chunk.append(sentence)
+            current_length += len(sentence)
+
+    if current_chunk:
+        chunks.append(" ".join(current_chunk))
+
+    # Process each chunk separately with Stanza
+    entities = []
+    for chunk in chunks:
+        doc = nlp(chunk)
+        for ent in doc.ents:
+            entities.append({"text": ent.text, "type": ent.type})
+
+    return entities
+
 ### 4️⃣ TTS Functionality (KokoroTTS)
 @spaces.GPU(duration=1000)  
 def generate_audio_kokoro(text, lang, selected_voice):
@@ -221,6 +256,7 @@ with gr.Blocks() as demo:
     process_text_button = gr.Button("Fetch Text & Detect Language")
     process_summary_button = gr.Button("Summarize Text", visible=False)
     process_audio_button = gr.Button("Generate Audio", visible=False)
+    process_ner_button = gr.Button("Extract Entities", visible=True)  # ✅ New button for NER
 
     # Layout: Two adjacent columns (Text and Metadata)
     with gr.Row():
@@ -233,6 +269,7 @@ with gr.Blocks() as demo:
     detected_lang = gr.Textbox(label="Detected Language", visible=False)
     summary_output = gr.Textbox(label="Summary", visible=True, interactive=False)
     full_audio_output = gr.Audio(label="Generated Audio", visible=True)
+    ner_output = gr.JSON(label="Extracted Entities", visible=True)  # ✅ New output for NER
 
     # Step 1: Fetch Text & Detect Language First
     process_text_button.click(
@@ -254,6 +291,12 @@ with gr.Blocks() as demo:
         outputs=[full_audio_output]
     )
 
+    process_ner_button.click(
+        extract_entities_with_stanza, 
+        inputs=[extracted_text], 
+        outputs=[ner_output]
+    )
+
     #process_button.click(process_url, inputs=[url_input], outputs=[extracted_text, detected_language, full_audio_output])
 
 demo.launch()