Create nlp_pipeline.py

nlp_pipeline.py

@@ -0,0 +1,66 @@
+# nlp_pipeline.py
+from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassification
+from sentence_transformers import SentenceTransformer
+import numpy as np
+
+# Load lighter/CPU-friendly models for HF Space
+SUMMARIZER = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=-1)
+# NER model (token-classification)
+NER = pipeline("ner", model="dbmdz/bert-large-cased-finetuned-conll03-english", aggregation_strategy="simple", device=-1)
+EMBED_MODEL = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")  # small & fast
+
+def summarize(text, max_length=120):
+    # chunk if needed
+    if len(text) < 800:
+        s = SUMMARIZER(text, max_length=max_length, min_length=40, do_sample=False)
+        return s[0]["summary_text"]
+    # naive chunking
+    parts = []
+    chunk_size = 700
+    for i in range(0, len(text), chunk_size):
+        chunk = text[i:i+chunk_size]
+        parts.append(SUMMARIZER(chunk, max_length=60, min_length=20)[0]["summary_text"])
+    return " ".join(parts)
+
+def extract_entities(text):
+    ner = NER(text)
+    # ner returns list of {'entity_group','score','word'}
+    grouped = {}
+    for ent in ner:
+        key = ent.get("entity_group") or ent.get("entity")
+        grouped.setdefault(key, []).append({"text": ent["word"], "score": float(ent["score"])})
+    return grouped
+
+def embed_text(text):
+    return EMBED_MODEL.encode(text, convert_to_numpy=True, normalize_embeddings=True)
+
+def get_sentence_provenance(sentences, entities):
+    # map entity text to sentences that contain it (case-insensitive)
+    prov = {}
+    for t in entities:
+        prov[t] = []
+        for s in sentences:
+            if t.lower() in s.lower():
+                prov[t].append(s)
+    return prov
+
+def process_document(doc):
+    text = doc["text"]
+    summary = summarize(text)
+    entities_grouped = extract_entities(text)
+    # flatten entity strings (unique)
+    entity_texts = set()
+    for k, v in entities_grouped.items():
+        for item in v:
+            entity_texts.add(item["text"])
+    provenance = get_sentence_provenance(doc["sentences"], entity_texts)
+    embedding = embed_text(summary)  # index the summary embedding for compactness
+    tags = []  # optional: simple tag by most frequent NER labels
+    return {
+        "summary": summary,
+        "entities": entities_grouped,
+        "entity_texts": list(entity_texts),
+        "provenance": provenance,
+        "embedding": embedding,
+        "tags": tags
+    }