SpacyBison

spacybison/config.cfg

@@ -0,0 +1,16 @@
+[nlp]
+lang = "en"
+pipeline = ["llm"]
+
+[components]
+
+[components.llm]
+factory = "llm"
+
+[components.llm.task]
+@llm_tasks = "spacy.NER.v2"
+labels = PERSON , FULL_ADDRESS , PHONE_NO , ORGANISATION , FIR_NO , CRIME_NO , DISTRICT , STATE , COUNTRY , AADHAAR , VEHICLE_NO
+
+[components.llm.model]
+@llm_models = "spacy.PaLM.v1"
+name = "text-bison-001"

spacybison/llm_ner.py

@@ -0,0 +1,10 @@
+import asyncio
+import json
+from spacy_llm.util import assemble
+
+def llm_ner(input_text):
+    nlp = assemble("spacybison/config.cfg")
+    doc = nlp(input_text)
+    entities = [(ent.text, ent.label_) for ent in doc.ents]
+
+    return entities

spacybison/occurences.py

@@ -0,0 +1,8 @@
+def find_all_occurrences(text, substring):
+    start = 0
+    while start < len(text):
+        start = text.find(substring, start)
+        if start == -1: 
+            break
+        yield start, start + len(substring)
+        start += len(substring)

spacybison/run_llm.py

@@ -0,0 +1,27 @@
+from spacybison.occurences import find_all_occurrences
+from spacybison.llm_ner import llm_ner
+import json
+
+def run_llm(input_text):
+    extracted_list = llm_ner(input_text)
+
+    output = []
+    seen_occurrences = set()
+
+    for entity, entity_type in extracted_list:
+        occurrences = list(find_all_occurrences(input_text, entity))
+        for start, end in occurrences:
+            if (start, end) not in seen_occurrences:
+                seen_occurrences.add((start, end))
+                output.append({
+                    "entity_type": entity_type,
+                    "start": start,
+                    "end": end,
+                    "score": 1,
+                    "analysis_explanation": None,
+                    "recognition_metadata": {
+                        "recognizer_identifier": "spacy-llm",
+                        "recognizer_name": "chat-bison-001"
+                    }
+                })
+    return output