structflo-ner 0.2.0__tar.gz → 0.2.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (37) hide show
  1. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/.github/workflows/ci.yml +1 -1
  2. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/PKG-INFO +4 -1
  3. structflo_ner-0.2.2/notebooks/02_fast_ner.ipynb +416 -0
  4. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/pyproject.toml +1 -0
  5. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/structflo/ner/__init__.py +9 -2
  6. structflo_ner-0.2.2/structflo/ner/_prompts.py +85 -0
  7. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/structflo/ner/extractor.py +63 -5
  8. structflo_ner-0.2.2/structflo/ner/fast/README.md +149 -0
  9. structflo_ner-0.2.2/structflo/ner/fast/__init__.py +22 -0
  10. structflo_ner-0.2.2/structflo/ner/fast/_loader.py +116 -0
  11. structflo_ner-0.2.2/structflo/ner/fast/_matcher.py +280 -0
  12. structflo_ner-0.2.2/structflo/ner/fast/_normalize.py +84 -0
  13. structflo_ner-0.2.2/structflo/ner/fast/extractor.py +160 -0
  14. structflo_ner-0.2.2/structflo/ner/fast/gazetteers/accession_number.yml +18 -0
  15. structflo_ner-0.2.2/structflo/ner/fast/gazetteers/compound_name.yml +50 -0
  16. structflo_ner-0.2.2/structflo/ner/fast/gazetteers/disease.yml +22 -0
  17. structflo_ner-0.2.2/structflo/ner/fast/gazetteers/functional_category.yml +27 -0
  18. structflo_ner-0.2.2/structflo/ner/fast/gazetteers/gene_name.yml +75 -0
  19. structflo_ner-0.2.2/structflo/ner/fast/gazetteers/product.yml +34 -0
  20. structflo_ner-0.2.2/structflo/ner/fast/gazetteers/screening_method.yml +38 -0
  21. structflo_ner-0.2.2/structflo/ner/fast/gazetteers/target.yml +78 -0
  22. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/tests/test_extractor.py +44 -0
  23. structflo_ner-0.2.2/tests/test_fast.py +377 -0
  24. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/uv.lock +97 -1
  25. structflo_ner-0.2.0/structflo/ner/_prompts.py +0 -133
  26. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/.github/workflows/publish.yml +0 -0
  27. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/.gitignore +0 -0
  28. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/Makefile +0 -0
  29. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/README.md +0 -0
  30. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/notebooks/01_quickstart.ipynb +0 -0
  31. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/structflo/ner/_display.py +0 -0
  32. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/structflo/ner/_entities.py +0 -0
  33. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/structflo/ner/_examples.py +5 -5
  34. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/structflo/ner/_mapping.py +0 -0
  35. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/structflo/ner/profiles.py +0 -0
  36. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/tests/__init__.py +0 -0
  37. {structflo_ner-0.2.0 → structflo_ner-0.2.2}/tests/test_entities.py +0 -0
@@ -19,7 +19,7 @@ jobs:
19
19
  run: uv python install 3.12
20
20
 
21
21
  - name: Install dependencies
22
- run: uv sync --dev
22
+ run: uv sync --dev --extra fast --extra dataframe
23
23
 
24
24
  - name: Lint
25
25
  run: uv run ruff check structflo/ tests/
@@ -1,12 +1,15 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: structflo-ner
3
- Version: 0.2.0
3
+ Version: 0.2.2
4
4
  Summary: Drug discovery NER wrapper around LangExtract — zero-config entity extraction for chemistry and biology.
5
5
  License: Apache-2.0
6
6
  Requires-Python: >=3.10
7
7
  Requires-Dist: langextract>=1.1.1
8
8
  Provides-Extra: dataframe
9
9
  Requires-Dist: pandas>=1.5; extra == 'dataframe'
10
+ Provides-Extra: fast
11
+ Requires-Dist: pyyaml>=6.0; extra == 'fast'
12
+ Requires-Dist: rapidfuzz>=3.0; extra == 'fast'
10
13
  Description-Content-Type: text/markdown
11
14
 
12
15
  # structflo.ner
@@ -0,0 +1,416 @@
1
+ {
2
+ "cells": [
3
+ {
4
+ "cell_type": "markdown",
5
+ "id": "cell-0",
6
+ "metadata": {},
7
+ "source": [
8
+ "# Fast Dictionary-Based NER for TB Drug Discovery\n",
9
+ "\n",
10
+ "This notebook demonstrates `FastNERExtractor` — a fast, deterministic NER\n",
11
+ "engine that uses curated YAML gazetteers instead of an LLM.\n",
12
+ "\n",
13
+ "1. Basic usage\n",
14
+ "2. How matching works (exact, regex, fuzzy)\n",
15
+ "3. Working with results (same as LLM extractor)\n",
16
+ "4. Custom gazetteers\n",
17
+ "5. Adding new entity types\n",
18
+ "6. Batch extraction & performance"
19
+ ]
20
+ },
21
+ {
22
+ "cell_type": "markdown",
23
+ "id": "cell-1",
24
+ "metadata": {},
25
+ "source": [
26
+ "## Setup\n",
27
+ "\n",
28
+ "```bash\n",
29
+ "uv add \"structflo-ner[fast]\"\n",
30
+ "# for DataFrame support\n",
31
+ "uv add \"structflo-ner[fast,dataframe]\"\n",
32
+ "```"
33
+ ]
34
+ },
35
+ {
36
+ "cell_type": "code",
37
+ "execution_count": null,
38
+ "id": "cell-2",
39
+ "metadata": {},
40
+ "outputs": [],
41
+ "source": [
42
+ "from structflo.ner.fast import FastNERExtractor"
43
+ ]
44
+ },
45
+ {
46
+ "cell_type": "markdown",
47
+ "id": "cell-3",
48
+ "metadata": {},
49
+ "source": [
50
+ "## 1. Basic Usage\n",
51
+ "\n",
52
+ "No API key, no LLM, no network — just instantiate and extract.\n",
53
+ "The built-in gazetteers cover TB drug discovery entities."
54
+ ]
55
+ },
56
+ {
57
+ "cell_type": "code",
58
+ "execution_count": null,
59
+ "id": "cell-4",
60
+ "metadata": {},
61
+ "outputs": [],
62
+ "source": [
63
+ "fast = FastNERExtractor()\n",
64
+ "\n",
65
+ "text = (\n",
66
+ " \"Bedaquiline (TMC207) is a diarylquinoline that inhibits the \"\n",
67
+ " \"mycobacterial ATP synthase subunit c encoded by atpE (Rv1305). \"\n",
68
+ " \"It shows potent activity against Mycobacterium tuberculosis \"\n",
69
+ " \"including MDR-TB and XDR-TB. The compound was identified through \"\n",
70
+ " \"whole-cell screening and targets the energy metabolism pathway.\"\n",
71
+ ")\n",
72
+ "\n",
73
+ "result = fast.extract(text)\n",
74
+ "result"
75
+ ]
76
+ },
77
+ {
78
+ "cell_type": "markdown",
79
+ "id": "cell-5",
80
+ "metadata": {},
81
+ "source": [
82
+ "## 2. How Matching Works\n",
83
+ "\n",
84
+ "The extractor uses a three-phase matching strategy:\n",
85
+ "\n",
86
+ "| Phase | Method | What it catches |\n",
87
+ "|---|---|---|\n",
88
+ "| 1 | **Exact match** | Case-sensitive and normalized dictionary lookups with word-boundary enforcement |\n",
89
+ "| 1b | **Regex patterns** | Auto-derived from accession number seeds (Rv tags, UniProt, PDB, etc.) |\n",
90
+ "| 2 | **Fuzzy match** | Typos and minor variants via rapidfuzz (configurable threshold) |\n",
91
+ "\n",
92
+ "Each entity's `attributes` dict includes the `match_method` used."
93
+ ]
94
+ },
95
+ {
96
+ "cell_type": "code",
97
+ "execution_count": null,
98
+ "id": "cell-6",
99
+ "metadata": {},
100
+ "outputs": [],
101
+ "source": [
102
+ "# Inspect match methods\n",
103
+ "for entity in result.all_entities():\n",
104
+ " method = entity.attributes.get(\"match_method\", \"\")\n",
105
+ " canonical = entity.attributes.get(\"canonical\", entity.text)\n",
106
+ " print(f\"{entity.entity_type:25s} | {entity.text:30s} | method={method:6s} | canonical={canonical}\")"
107
+ ]
108
+ },
109
+ {
110
+ "cell_type": "markdown",
111
+ "id": "cell-7",
112
+ "metadata": {},
113
+ "source": [
114
+ "### Regex matching for accession numbers\n",
115
+ "\n",
116
+ "Seed entries in `accession_number.yml` auto-derive regex patterns.\n",
117
+ "For example, `Rv0005` teaches the system `Rv\\d{4}[c]?`, so *all* Rv locus tags\n",
118
+ "are matched — not just the ones listed."
119
+ ]
120
+ },
121
+ {
122
+ "cell_type": "code",
123
+ "execution_count": null,
124
+ "id": "cell-8",
125
+ "metadata": {},
126
+ "outputs": [],
127
+ "source": [
128
+ "# Rv2043c is NOT in the YAML, but the regex pattern catches it\n",
129
+ "result2 = fast.extract(\"PptT is encoded by Rv2043c and is essential for mycolic acid biosynthesis.\")\n",
130
+ "\n",
131
+ "print(\"Accessions found:\")\n",
132
+ "for a in result2.accessions:\n",
133
+ " print(f\" {a.text} (method: {a.attributes['match_method']})\")\n",
134
+ "\n",
135
+ "print(\"\\nTargets found:\")\n",
136
+ "for t in result2.targets:\n",
137
+ " print(f\" {t.text}\")"
138
+ ]
139
+ },
140
+ {
141
+ "cell_type": "markdown",
142
+ "id": "cell-9",
143
+ "metadata": {},
144
+ "source": [
145
+ "### Fuzzy matching\n",
146
+ "\n",
147
+ "Catches typos and minor spelling variants. The threshold (default 85) controls sensitivity."
148
+ ]
149
+ },
150
+ {
151
+ "cell_type": "code",
152
+ "execution_count": null,
153
+ "id": "cell-10",
154
+ "metadata": {},
155
+ "outputs": [],
156
+ "source": [
157
+ "# \"Bedaquilne\" is a typo for \"Bedaquiline\"\n",
158
+ "fuzzy_result = fast.extract(\"Bedaquilne showed activity against TB\")\n",
159
+ "\n",
160
+ "for c in fuzzy_result.compounds:\n",
161
+ " print(f\"Found: {c.text!r} -> canonical: {c.attributes.get('canonical', c.text)!r} (method: {c.attributes['match_method']})\")"
162
+ ]
163
+ },
164
+ {
165
+ "cell_type": "code",
166
+ "execution_count": null,
167
+ "id": "cell-11",
168
+ "metadata": {},
169
+ "outputs": [],
170
+ "source": [
171
+ "# Disable fuzzy matching for strict mode\n",
172
+ "strict = FastNERExtractor(fuzzy_threshold=0)\n",
173
+ "strict_result = strict.extract(\"Bedaquilne showed activity against TB\")\n",
174
+ "\n",
175
+ "print(f\"Compounds (strict): {strict_result.compounds}\") # empty — typo not matched\n",
176
+ "print(f\"Diseases (strict): {[d.text for d in strict_result.diseases]}\") # TB still matched"
177
+ ]
178
+ },
179
+ {
180
+ "cell_type": "markdown",
181
+ "id": "cell-12",
182
+ "metadata": {},
183
+ "source": [
184
+ "## 3. Working with Results\n",
185
+ "\n",
186
+ "`FastNERExtractor` returns the same `NERResult` objects as the LLM-based `NERExtractor`.\n",
187
+ "All downstream tooling works identically."
188
+ ]
189
+ },
190
+ {
191
+ "cell_type": "code",
192
+ "execution_count": null,
193
+ "id": "cell-13",
194
+ "metadata": {},
195
+ "outputs": [],
196
+ "source": [
197
+ "# Typed entity lists\n",
198
+ "print(\"Compounds:\", [c.text for c in result.compounds])\n",
199
+ "print(\"Targets:\", [t.text for t in result.targets])\n",
200
+ "print(\"Diseases:\", [d.text for d in result.diseases])\n",
201
+ "print(\"Accessions:\", [a.text for a in result.accessions])\n",
202
+ "print(\"Screening methods:\", [s.text for s in result.screening_methods])\n",
203
+ "print(\"Products:\", [p.text for p in result.products])\n",
204
+ "print(\"Functional categories:\", [f.text for f in result.functional_categories])"
205
+ ]
206
+ },
207
+ {
208
+ "cell_type": "code",
209
+ "execution_count": null,
210
+ "id": "cell-14",
211
+ "metadata": {},
212
+ "outputs": [],
213
+ "source": [
214
+ "# Export to DataFrame\n",
215
+ "df = result.to_dataframe()\n",
216
+ "df"
217
+ ]
218
+ },
219
+ {
220
+ "cell_type": "code",
221
+ "execution_count": null,
222
+ "id": "cell-15",
223
+ "metadata": {},
224
+ "outputs": [],
225
+ "source": [
226
+ "# Character offsets let you highlight entities in the source text\n",
227
+ "for entity in result.all_entities():\n",
228
+ " if entity.char_start is not None:\n",
229
+ " span = text[entity.char_start:entity.char_end]\n",
230
+ " print(f\"[{entity.char_start:3d}:{entity.char_end:3d}] {entity.entity_type:25s} | {span!r}\")"
231
+ ]
232
+ },
233
+ {
234
+ "cell_type": "markdown",
235
+ "id": "cell-16",
236
+ "metadata": {},
237
+ "source": [
238
+ "## 4. Custom Gazetteers\n",
239
+ "\n",
240
+ "YAML files contain only names — one per line. The filename becomes the `entity_type`.\n",
241
+ "\n",
242
+ "```yaml\n",
243
+ "# my_gazetteers/target.yml\n",
244
+ "- MyNovelTarget\n",
245
+ "- AnotherTarget\n",
246
+ "- KinaseX\n",
247
+ "```\n",
248
+ "\n",
249
+ "```python\n",
250
+ "fast = FastNERExtractor(gazetteer_dir=\"my_gazetteers/\")\n",
251
+ "```\n",
252
+ "\n",
253
+ "You can also add terms programmatically without creating files:"
254
+ ]
255
+ },
256
+ {
257
+ "cell_type": "code",
258
+ "execution_count": null,
259
+ "id": "cell-17",
260
+ "metadata": {},
261
+ "outputs": [],
262
+ "source": [
263
+ "# Add extra terms on top of the built-in gazetteers\n",
264
+ "custom = FastNERExtractor(\n",
265
+ " extra_gazetteers={\n",
266
+ " \"target\": [\"MyNovelTarget\", \"KinaseX\"],\n",
267
+ " \"compound_name\": [\"CompoundABC\"],\n",
268
+ " }\n",
269
+ ")\n",
270
+ "\n",
271
+ "r = custom.extract(\"CompoundABC inhibits MyNovelTarget in M. tuberculosis\")\n",
272
+ "print(\"Compounds:\", [c.text for c in r.compounds])\n",
273
+ "print(\"Targets:\", [t.text for t in r.targets])\n",
274
+ "print(\"Diseases:\", [d.text for d in r.diseases])"
275
+ ]
276
+ },
277
+ {
278
+ "cell_type": "markdown",
279
+ "id": "cell-18",
280
+ "metadata": {},
281
+ "source": [
282
+ "## 5. Adding New Entity Types\n",
283
+ "\n",
284
+ "To add a new gazetteer, just drop a YAML file into the gazetteers directory.\n",
285
+ "The filename (without `.yml`) must match an `entity_type` from the entity class map.\n",
286
+ "\n",
287
+ "**Built-in entity types:**\n",
288
+ "\n",
289
+ "| Filename | entity_type | Python class |\n",
290
+ "|---|---|---|\n",
291
+ "| `target.yml` | target | `TargetEntity` |\n",
292
+ "| `gene_name.yml` | gene_name | `TargetEntity` |\n",
293
+ "| `compound_name.yml` | compound_name | `ChemicalEntity` |\n",
294
+ "| `disease.yml` | disease | `DiseaseEntity` |\n",
295
+ "| `accession_number.yml` | accession_number | `AccessionEntity` |\n",
296
+ "| `screening_method.yml` | screening_method | `ScreeningMethodEntity` |\n",
297
+ "| `functional_category.yml` | functional_category | `FunctionalCategoryEntity` |\n",
298
+ "| `product.yml` | product | `ProductEntity` |\n",
299
+ "\n",
300
+ "**What's auto-derived from names:**\n",
301
+ "- Case variants (InhA, inha, INHA)\n",
302
+ "- Hyphen-optional forms (DprE-1 ↔ DprE1)\n",
303
+ "- Period-optional forms (M. tuberculosis ↔ M tuberculosis)\n",
304
+ "- Greek letter expansion (β-lactam ↔ beta-lactam)\n",
305
+ "- Regex patterns for accession number seeds (Rv, MT, UniProt, PDB, RefSeq)"
306
+ ]
307
+ },
308
+ {
309
+ "cell_type": "code",
310
+ "execution_count": null,
311
+ "id": "cell-19",
312
+ "metadata": {},
313
+ "outputs": [],
314
+ "source": [
315
+ "# See what gazetteers are loaded by default\n",
316
+ "from structflo.ner.fast._loader import load_all_gazetteers\n",
317
+ "\n",
318
+ "gazetteers = load_all_gazetteers()\n",
319
+ "for entity_type, terms in gazetteers.items():\n",
320
+ " print(f\"{entity_type:25s} | {len(terms):3d} terms | first 5: {terms[:5]}\")"
321
+ ]
322
+ },
323
+ {
324
+ "cell_type": "markdown",
325
+ "id": "cell-20",
326
+ "metadata": {},
327
+ "source": [
328
+ "## 6. Batch Extraction & Performance\n",
329
+ "\n",
330
+ "The fast extractor processes text in milliseconds — orders of magnitude faster than LLM-based extraction."
331
+ ]
332
+ },
333
+ {
334
+ "cell_type": "code",
335
+ "execution_count": null,
336
+ "id": "cell-21",
337
+ "metadata": {},
338
+ "outputs": [],
339
+ "source": [
340
+ "abstracts = [\n",
341
+ " \"Bedaquiline inhibits AtpE (Rv1305) with nanomolar potency against MDR-TB.\",\n",
342
+ " \"Delamanid (OPC-67683) is activated by Ddn and targets mycolic acid biosynthesis in M. tuberculosis.\",\n",
343
+ " \"Pretomanid (PA-824) requires activation by Ddn (Rv3547) and kills both replicating and non-replicating Mtb.\",\n",
344
+ " \"PBTZ169 (Macozinone) inhibits DprE1 (Rv3790), an essential enzyme in cell wall biosynthesis.\",\n",
345
+ " \"SQ109 targets MmpL3, a trehalose monomycolate transporter essential for cell wall assembly.\",\n",
346
+ " \"Fragment-based screening identified InhA inhibitors that bypass katG-mediated activation.\",\n",
347
+ " \"CRISPRi screening revealed QcrB (Rv2196) as a vulnerable target in energy metabolism.\",\n",
348
+ " \"Structure-based drug design targeting KasA (Rv2245) yielded novel fatty acid biosynthesis inhibitors.\",\n",
349
+ "]\n",
350
+ "\n",
351
+ "results = fast.extract(abstracts)\n",
352
+ "\n",
353
+ "for i, r in enumerate(results):\n",
354
+ " entities = r.all_entities()\n",
355
+ " print(f\"Abstract {i+1}: {len(entities)} entities — {[e.text for e in entities]}\")"
356
+ ]
357
+ },
358
+ {
359
+ "cell_type": "code",
360
+ "execution_count": null,
361
+ "id": "cell-22",
362
+ "metadata": {},
363
+ "outputs": [],
364
+ "source": [
365
+ "%%timeit -n 100\n",
366
+ "# Benchmark: extract from a single abstract\n",
367
+ "fast.extract(text)"
368
+ ]
369
+ },
370
+ {
371
+ "cell_type": "code",
372
+ "execution_count": null,
373
+ "id": "cell-23",
374
+ "metadata": {},
375
+ "outputs": [],
376
+ "source": [
377
+ "%%timeit -n 10\n",
378
+ "# Benchmark: extract from 8 abstracts\n",
379
+ "fast.extract(abstracts)"
380
+ ]
381
+ },
382
+ {
383
+ "cell_type": "markdown",
384
+ "id": "cell-24",
385
+ "metadata": {},
386
+ "source": [
387
+ "## Comparing Fast vs LLM Extraction\n",
388
+ "\n",
389
+ "The fast extractor is ideal as a **first pass** for bulk screening.\n",
390
+ "Use the LLM extractor for deeper analysis where context and novel entities matter.\n",
391
+ "\n",
392
+ "| | `FastNERExtractor` | `NERExtractor` |\n",
393
+ "|---|---|---|\n",
394
+ "| Speed | ~1-5ms per abstract | ~2-5s per abstract |\n",
395
+ "| Novel entities | Only known terms | Discovers new entities |\n",
396
+ "| Context awareness | None (string matching) | Full contextual understanding |\n",
397
+ "| Cost | Free (no API calls) | API costs or GPU |\n",
398
+ "| Setup | Zero config | API key or Ollama |\n",
399
+ "| Output | `NERResult` | `NERResult` (identical) |"
400
+ ]
401
+ }
402
+ ],
403
+ "metadata": {
404
+ "kernelspec": {
405
+ "display_name": "Python 3",
406
+ "language": "python",
407
+ "name": "python3"
408
+ },
409
+ "language_info": {
410
+ "name": "python",
411
+ "version": "3.12.0"
412
+ }
413
+ },
414
+ "nbformat": 4,
415
+ "nbformat_minor": 5
416
+ }
@@ -11,6 +11,7 @@ dependencies = [
11
11
 
12
12
  [project.optional-dependencies]
13
13
  dataframe = ["pandas>=1.5"]
14
+ fast = ["rapidfuzz>=3.0", "PyYAML>=6.0"]
14
15
 
15
16
  [dependency-groups]
16
17
  dev = [
@@ -49,6 +49,12 @@ from structflo.ner._entities import (
49
49
  TargetEntity,
50
50
  )
51
51
  from structflo.ner.extractor import NERExtractor
52
+
53
+ try:
54
+ from structflo.ner.fast import FastNERExtractor
55
+ except ImportError: # rapidfuzz / PyYAML not installed
56
+ FastNERExtractor = None # type: ignore[assignment,misc]
57
+
52
58
  from structflo.ner.profiles import (
53
59
  BIOACTIVITY,
54
60
  BIOLOGY,
@@ -61,11 +67,12 @@ from structflo.ner.profiles import (
61
67
  EntityProfile,
62
68
  )
63
69
 
64
- __version__ = "0.2.0"
70
+ __version__ = "0.2.2"
65
71
 
66
72
  __all__ = [
67
- # Main class
73
+ # Main classes
68
74
  "NERExtractor",
75
+ "FastNERExtractor",
69
76
  # Profile system
70
77
  "EntityProfile",
71
78
  "FULL",
@@ -0,0 +1,85 @@
1
+ """Prompt strings for each built-in EntityProfile."""
2
+
3
+ CHEMISTRY_PROMPT = (
4
+ "Extract chemical entities from the text. "
5
+ "Include: compound names (generic names, IUPAC names, code names like 'Compound 5b', "
6
+ "brand names, and abbreviations), SMILES strings (the exact SMILES notation as written), "
7
+ "CAS registry numbers (e.g. '50-78-2'), and molecular formulas (e.g. 'C9H8O4'). "
8
+ "Do not infer or generate SMILES — only extract them if explicitly present in the text."
9
+ )
10
+
11
+ BIOLOGY_PROMPT = (
12
+ "Extract biological target entities from the text. "
13
+ "Include: protein targets (e.g. 'EGFR', 'CDK4/6', 'PD-L1'), gene names (e.g. 'KRAS', 'TP53', "
14
+ "'BRCA1'), receptor names, enzyme names, and pathway names. "
15
+ "For each target, capture the gene symbol if mentioned alongside a protein name. "
16
+ "Capture the organism context if specified (e.g. 'human', 'mouse')."
17
+ )
18
+
19
+ BIOACTIVITY_PROMPT = (
20
+ "Extract bioactivity measurements and assay data from the text. "
21
+ "Include: potency values (IC50, EC50, Ki, Kd, GI50, CC50), selectivity ratios, "
22
+ "percent inhibition values, and Hill coefficients. "
23
+ "For each value, capture the numeric value, unit (nM, µM, mM), and measurement type. "
24
+ "Also extract assay descriptions: cell lines used (e.g. 'HeLa', 'A549'), assay formats "
25
+ "(e.g. 'cell viability', 'binding assay', 'enzymatic assay'), and organisms."
26
+ )
27
+
28
+ DISEASE_PROMPT = (
29
+ "Extract disease names and clinical indications from the text. "
30
+ "Include: cancer types (e.g. 'non-small cell lung cancer', 'AML', 'NSCLC'), "
31
+ "non-oncology diseases (e.g. 'type 2 diabetes', 'rheumatoid arthritis'), "
32
+ "and therapeutic areas (e.g. 'oncology', 'CNS'). "
33
+ "Capture both full names and abbreviations. "
34
+ "For each disease, note the therapeutic area if discernible from context."
35
+ )
36
+
37
+ FULL_PROMPT = (
38
+ "Extract all drug discovery entities from the text. "
39
+ "This includes:\n"
40
+ "- Chemical entities: compound names (generic, IUPAC, code names, brand names), "
41
+ "SMILES strings (only if explicitly written), CAS numbers, molecular formulas.\n"
42
+ "- Biological targets: protein names, gene names, receptor names, enzyme names, pathways.\n"
43
+ "- Bioactivity data: IC50, EC50, Ki, Kd, and other potency/selectivity measurements "
44
+ "with their numeric values and units.\n"
45
+ "- Assay information: cell lines, assay formats, experimental organisms.\n"
46
+ "- Diseases and indications: cancer types, disease names, therapeutic areas.\n"
47
+ "- Mechanisms of action: binding modes, inhibition types, selectivity descriptions.\n"
48
+ "Extract only what is explicitly stated; do not infer or generate values."
49
+ )
50
+
51
+ # ── Tuberculosis early drug discovery prompts ──────────────────────────
52
+
53
+ TB_PROMPT = (
54
+ "Extract drug discovery entities from this tuberculosis research text.\n\n"
55
+ "DISAMBIGUATION RULES:\n"
56
+ "- Mycobacterial proteins (e.g. ClpC1, DprE1, InhA, AtpE, MmpL3, QcrB) "
57
+ "are biological targets, NOT compounds.\n"
58
+ "- Rv locus tags (Rv3790, Rv1484), UniProt IDs (P9WPS1), and PDB codes "
59
+ "are accession_number, not target or gene_name.\n"
60
+ "- Enzyme descriptions like 'enoyl-ACP reductase' are product, not target.\n"
61
+ "- 'cell wall', 'lipid metabolism' are functional_category, not mechanism_of_action.\n"
62
+ "- 'fragment screening', 'biochemical assay' are screening_method, not assay.\n"
63
+ "- Use target for proteins in a drug-targeting context, gene_name for loci, "
64
+ "protein_name for non-drug-target proteins.\n\n"
65
+ "Extract only what is explicitly stated; do not infer or generate values."
66
+ )
67
+
68
+ TB_CHEMISTRY_PROMPT = (
69
+ "Extract chemical entities from this tuberculosis drug discovery text. "
70
+ "Include compound names, SMILES (only if explicitly written), CAS numbers, "
71
+ "and molecular formulas. "
72
+ "Mycobacterial proteins (ClpC1, DprE1, InhA, AtpE, etc.) are NOT compounds. "
73
+ "Extract only what is explicitly stated; do not infer or generate values."
74
+ )
75
+
76
+ TB_BIOLOGY_PROMPT = (
77
+ "Extract biological entities from this tuberculosis research text. "
78
+ "Use target for proteins in a drug-targeting context, gene_name for loci, "
79
+ "protein_name for non-drug-target proteins. "
80
+ "Rv locus tags and UniProt IDs are accession_number. "
81
+ "Enzyme descriptions (e.g. 'enoyl-ACP reductase') are product. "
82
+ "Protein functional categories (e.g. 'cell wall', 'lipid metabolism') "
83
+ "are functional_category. "
84
+ "Extract only what is explicitly stated; do not infer or generate values."
85
+ )
@@ -2,12 +2,16 @@
2
2
 
3
3
  from __future__ import annotations
4
4
 
5
+ import logging
6
+
5
7
  import langextract as lx
6
8
 
7
9
  from structflo.ner._entities import NERResult
8
10
  from structflo.ner._mapping import annotated_doc_to_result
9
11
  from structflo.ner.profiles import FULL, EntityProfile
10
12
 
13
+ logger = logging.getLogger(__name__)
14
+
11
15
 
12
16
  class NERExtractor:
13
17
  """Extract drug discovery entities from text with zero configuration.
@@ -109,8 +113,53 @@ class NERExtractor:
109
113
  return profile.examples + self._extra_examples
110
114
 
111
115
  def _build_prompt(self, profile: EntityProfile) -> str:
112
- """Return the prompt string for the given profile."""
113
- return profile.prompt
116
+ """Return the prompt string for the given profile.
117
+
118
+ Appends an explicit schema constraint listing the allowed
119
+ entity classes. This is critical for models that don't support
120
+ structured-output schemas (e.g. Ollama) where
121
+ ``use_schema_constraints`` has no effect.
122
+ """
123
+ classes = ", ".join(profile.entity_classes)
124
+ constraint = (
125
+ f"\n\nIMPORTANT — You MUST classify every extraction using "
126
+ f"ONLY these entity classes: [{classes}]. "
127
+ f"Do NOT invent new class names. Any extraction_class not in "
128
+ f"this list is an error."
129
+ )
130
+ return profile.prompt + constraint
131
+
132
+ @staticmethod
133
+ def _filter_extractions(
134
+ doc: lx.data.AnnotatedDocument,
135
+ allowed_classes: set[str],
136
+ ) -> lx.data.AnnotatedDocument:
137
+ """Drop extractions whose class is not in the profile's allowed set.
138
+
139
+ Models without schema-constraint support (e.g. Ollama) may invent
140
+ arbitrary entity classes. This post-processing step removes those
141
+ hallucinated classes so they don't pollute the result.
142
+ """
143
+ kept: list[lx.data.Extraction] = []
144
+ for ext in doc.extractions:
145
+ if ext.extraction_class in allowed_classes:
146
+ kept.append(ext)
147
+ else:
148
+ logger.warning(
149
+ "Dropping extraction with unknown class %r (text=%r). Allowed classes: %s",
150
+ ext.extraction_class,
151
+ ext.extraction_text,
152
+ ", ".join(sorted(allowed_classes)),
153
+ )
154
+ return lx.data.AnnotatedDocument(
155
+ text=doc.text,
156
+ extractions=kept,
157
+ )
158
+
159
+ @property
160
+ def _is_ollama(self) -> bool:
161
+ """Return True when routing to an Ollama endpoint."""
162
+ return self._model_url is not None
114
163
 
115
164
  def _run_extraction(
116
165
  self,
@@ -125,6 +174,13 @@ class NERExtractor:
125
174
  kwargs.setdefault("use_schema_constraints", True)
126
175
  kwargs.setdefault("show_progress", False)
127
176
 
177
+ # Ollama defaults to num_ctx=2048 which is far too small for
178
+ # few-shot NER prompts. Set a sane default so users don't hit
179
+ # silent truncation.
180
+ if self._is_ollama:
181
+ lm_params = kwargs.setdefault("language_model_params", {})
182
+ lm_params.setdefault("num_ctx", 8192)
183
+
128
184
  result = lx.extract(
129
185
  text_or_documents=text,
130
186
  prompt_description=prompt,
@@ -136,6 +192,8 @@ class NERExtractor:
136
192
  )
137
193
 
138
194
  # lx.extract returns a list when given a list; we always pass a single string
139
- if isinstance(result, list):
140
- return result[0]
141
- return result
195
+ doc = result[0] if isinstance(result, list) else result
196
+
197
+ # Post-process: drop extractions with hallucinated entity classes
198
+ allowed = set(profile.entity_classes)
199
+ return self._filter_extractions(doc, allowed)