schematize 0.1.0__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 (76) hide show
  1. schematize-0.1.0/.gitignore +52 -0
  2. schematize-0.1.0/LICENSE +21 -0
  3. schematize-0.1.0/PKG-INFO +307 -0
  4. schematize-0.1.0/README.md +259 -0
  5. schematize-0.1.0/pyproject.toml +106 -0
  6. schematize-0.1.0/src/schematize/__init__.py +32 -0
  7. schematize-0.1.0/src/schematize/_scripts/__init__.py +0 -0
  8. schematize-0.1.0/src/schematize/_scripts/evaluate_annotator_agreement.py +100 -0
  9. schematize-0.1.0/src/schematize/_scripts/evaluate_schema.py +169 -0
  10. schematize-0.1.0/src/schematize/_scripts/schema_generator.py +116 -0
  11. schematize-0.1.0/src/schematize/_scripts/schema_generator_ablation.py +151 -0
  12. schematize-0.1.0/src/schematize/_scripts/schema_generator_mocked.py +153 -0
  13. schematize-0.1.0/src/schematize/agents/__init__.py +0 -0
  14. schematize-0.1.0/src/schematize/agents/agent_state.py +54 -0
  15. schematize-0.1.0/src/schematize/agents/basic_agents.py +115 -0
  16. schematize-0.1.0/src/schematize/agents/chat_agent.py +85 -0
  17. schematize-0.1.0/src/schematize/agents/data_agents.py +165 -0
  18. schematize-0.1.0/src/schematize/agents/output_models.py +50 -0
  19. schematize-0.1.0/src/schematize/agents/schema_generator.py +339 -0
  20. schematize-0.1.0/src/schematize/configs/__init__.py +0 -0
  21. schematize-0.1.0/src/schematize/configs/prompt/en/law/init_chat.yaml +137 -0
  22. schematize-0.1.0/src/schematize/configs/prompt/en/law/problem_definer.yaml +52 -0
  23. schematize-0.1.0/src/schematize/configs/prompt/en/law/problem_definer_helper.yaml +119 -0
  24. schematize-0.1.0/src/schematize/configs/prompt/en/law/query_generator.yaml +32 -0
  25. schematize-0.1.0/src/schematize/configs/prompt/en/law/schema_assessment.yaml +77 -0
  26. schematize-0.1.0/src/schematize/configs/prompt/en/law/schema_data_assessment.yaml +71 -0
  27. schematize-0.1.0/src/schematize/configs/prompt/en/law/schema_data_assessment_merger.yaml +93 -0
  28. schematize-0.1.0/src/schematize/configs/prompt/en/law/schema_data_refiner.yaml +125 -0
  29. schematize-0.1.0/src/schematize/configs/prompt/en/law/schema_generator.yaml +65 -0
  30. schematize-0.1.0/src/schematize/configs/prompt/en/law/schema_refiner.yaml +88 -0
  31. schematize-0.1.0/src/schematize/configs/prompt/en/tax/init_chat.yaml +128 -0
  32. schematize-0.1.0/src/schematize/configs/prompt/en/tax/problem_definer.yaml +64 -0
  33. schematize-0.1.0/src/schematize/configs/prompt/en/tax/problem_definer_helper.yaml +100 -0
  34. schematize-0.1.0/src/schematize/configs/prompt/en/tax/query_generator.yaml +40 -0
  35. schematize-0.1.0/src/schematize/configs/prompt/en/tax/schema_assessment.yaml +78 -0
  36. schematize-0.1.0/src/schematize/configs/prompt/en/tax/schema_data_assessment.yaml +72 -0
  37. schematize-0.1.0/src/schematize/configs/prompt/en/tax/schema_data_assessment_merger.yaml +94 -0
  38. schematize-0.1.0/src/schematize/configs/prompt/en/tax/schema_data_refiner.yaml +126 -0
  39. schematize-0.1.0/src/schematize/configs/prompt/en/tax/schema_generator.yaml +66 -0
  40. schematize-0.1.0/src/schematize/configs/prompt/en/tax/schema_refiner.yaml +90 -0
  41. schematize-0.1.0/src/schematize/configs/prompt/eval/annotator_agreement.yaml +31 -0
  42. schematize-0.1.0/src/schematize/configs/prompt/eval/schema_evaluator.yaml +40 -0
  43. schematize-0.1.0/src/schematize/configs/prompt/pl/law/init_chat.yaml +133 -0
  44. schematize-0.1.0/src/schematize/configs/prompt/pl/law/problem_definer.yaml +52 -0
  45. schematize-0.1.0/src/schematize/configs/prompt/pl/law/problem_definer_helper.yaml +158 -0
  46. schematize-0.1.0/src/schematize/configs/prompt/pl/law/query_generator.yaml +32 -0
  47. schematize-0.1.0/src/schematize/configs/prompt/pl/law/schema_assessment.yaml +67 -0
  48. schematize-0.1.0/src/schematize/configs/prompt/pl/law/schema_data_assessment.yaml +70 -0
  49. schematize-0.1.0/src/schematize/configs/prompt/pl/law/schema_data_assessment_merger.yaml +90 -0
  50. schematize-0.1.0/src/schematize/configs/prompt/pl/law/schema_data_refiner.yaml +120 -0
  51. schematize-0.1.0/src/schematize/configs/prompt/pl/law/schema_generator.yaml +65 -0
  52. schematize-0.1.0/src/schematize/configs/prompt/pl/law/schema_refiner.yaml +88 -0
  53. schematize-0.1.0/src/schematize/configs/prompt/pl/tax/init_chat.yaml +126 -0
  54. schematize-0.1.0/src/schematize/configs/prompt/pl/tax/problem_definer.yaml +62 -0
  55. schematize-0.1.0/src/schematize/configs/prompt/pl/tax/problem_definer_helper.yaml +156 -0
  56. schematize-0.1.0/src/schematize/configs/prompt/pl/tax/query_generator.yaml +40 -0
  57. schematize-0.1.0/src/schematize/configs/prompt/pl/tax/schema_assessment.yaml +68 -0
  58. schematize-0.1.0/src/schematize/configs/prompt/pl/tax/schema_data_assessment.yaml +71 -0
  59. schematize-0.1.0/src/schematize/configs/prompt/pl/tax/schema_data_assessment_merger.yaml +93 -0
  60. schematize-0.1.0/src/schematize/configs/prompt/pl/tax/schema_data_refiner.yaml +121 -0
  61. schematize-0.1.0/src/schematize/configs/prompt/pl/tax/schema_generator.yaml +61 -0
  62. schematize-0.1.0/src/schematize/configs/prompt/pl/tax/schema_refiner.yaml +90 -0
  63. schematize-0.1.0/src/schematize/eval/__init__.py +0 -0
  64. schematize-0.1.0/src/schematize/eval/evaluator.py +218 -0
  65. schematize-0.1.0/src/schematize/py.typed +0 -0
  66. schematize-0.1.0/src/schematize/retrieval/__init__.py +3 -0
  67. schematize-0.1.0/src/schematize/retrieval/base.py +13 -0
  68. schematize-0.1.0/src/schematize/retrieval/huggingface.py +232 -0
  69. schematize-0.1.0/src/schematize/retrieval/weaviate.py +125 -0
  70. schematize-0.1.0/src/schematize/schema/__init__.py +0 -0
  71. schematize-0.1.0/src/schematize/schema/model.py +74 -0
  72. schematize-0.1.0/src/schematize/settings.py +8 -0
  73. schematize-0.1.0/src/schematize/utils/__init__.py +0 -0
  74. schematize-0.1.0/src/schematize/utils/langchain.py +10 -0
  75. schematize-0.1.0/src/schematize/utils/load.py +25 -0
  76. schematize-0.1.0/src/schematize/utils/retry.py +70 -0
@@ -0,0 +1,52 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ *.so
4
+ *.egg
5
+ *.egg-info/
6
+ dist/
7
+ build/
8
+ .eggs/
9
+ .venv/
10
+ venv/
11
+ .uv/
12
+
13
+ # hydra outputs
14
+ outputs/
15
+ multirun/
16
+
17
+ # langchain cache
18
+ *.langchain_cache.db
19
+ .langchain_cache.db
20
+
21
+ # env
22
+ .env
23
+ *.env
24
+
25
+ # testing
26
+ .coverage
27
+ htmlcov/
28
+ .pytest_cache/
29
+
30
+ # mypy
31
+ .mypy_cache/
32
+
33
+ # ruff
34
+ .ruff_cache/
35
+
36
+ # editors
37
+ .idea/
38
+ .vscode/
39
+ *.swp
40
+
41
+ # notebooks
42
+ .ipynb_checkpoints/
43
+
44
+ # OS
45
+ .DS_Store
46
+
47
+ # baselines: third-party clone + local retrieval index cache
48
+ baselines/ScheMatiQ/
49
+ .cache/
50
+
51
+ # mkdocs build output
52
+ site/
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 schematize contributors
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,307 @@
1
+ Metadata-Version: 2.4
2
+ Name: schematize
3
+ Version: 0.1.0
4
+ Summary: Agentic system for automated extraction-schema generation from natural language descriptions
5
+ Author-email: Albert Sawczyn <albert.sawczyn@pwr.edu.pl>, Jakub Binkowski <jakub.binkowski@pwr.edu.pl>
6
+ License: MIT
7
+ License-File: LICENSE
8
+ Requires-Python: >=3.12
9
+ Requires-Dist: langchain-community>=0.3
10
+ Requires-Dist: langchain-core>=0.3
11
+ Requires-Dist: langchain-openai>=0.3
12
+ Requires-Dist: langchain-text-splitters>=0.3
13
+ Requires-Dist: langchain>=0.3
14
+ Requires-Dist: langgraph>=0.6
15
+ Requires-Dist: loguru>=0.7
16
+ Requires-Dist: pydantic>=2.0
17
+ Requires-Dist: python-dotenv>=1.0.0
18
+ Requires-Dist: pyyaml>=6.0
19
+ Requires-Dist: typer>=0.26.7
20
+ Requires-Dist: typing-extensions>=4.0
21
+ Provides-Extra: dev
22
+ Requires-Dist: coverage>=7.0; extra == 'dev'
23
+ Requires-Dist: mypy>=1.0; extra == 'dev'
24
+ Requires-Dist: pre-commit>=4.0; extra == 'dev'
25
+ Requires-Dist: pytest>=8.0; extra == 'dev'
26
+ Requires-Dist: ruff>=0.1; extra == 'dev'
27
+ Provides-Extra: docs
28
+ Requires-Dist: mkdocs-material>=9.5; extra == 'docs'
29
+ Requires-Dist: mkdocs>=1.6; extra == 'docs'
30
+ Requires-Dist: mkdocstrings[python]>=0.25; extra == 'docs'
31
+ Provides-Extra: huggingface
32
+ Requires-Dist: datasets>=3.0; extra == 'huggingface'
33
+ Requires-Dist: faiss-cpu>=1.8; extra == 'huggingface'
34
+ Requires-Dist: sentence-transformers>=3.0; extra == 'huggingface'
35
+ Provides-Extra: notebooks
36
+ Requires-Dist: ipykernel>=6.0; extra == 'notebooks'
37
+ Requires-Dist: jupyter>=1.0; extra == 'notebooks'
38
+ Requires-Dist: matplotlib>=3.0; extra == 'notebooks'
39
+ Requires-Dist: pandas>=2.0; extra == 'notebooks'
40
+ Requires-Dist: seaborn>=0.13; extra == 'notebooks'
41
+ Provides-Extra: scripts
42
+ Requires-Dist: hydra-core>=1.3; extra == 'scripts'
43
+ Requires-Dist: hydra-joblib-launcher>=1.2; extra == 'scripts'
44
+ Requires-Dist: tqdm>=4.0; extra == 'scripts'
45
+ Provides-Extra: weaviate
46
+ Requires-Dist: weaviate-client>=4.11; extra == 'weaviate'
47
+ Description-Content-Type: text/markdown
48
+
49
+ <div align="center">
50
+
51
+ <img src="docs/assets/logo.png" alt="schematize logo" width="200" height="200">
52
+
53
+ # schematize
54
+
55
+ **Turn a plain-English research problem into a typed, data-tested extraction schema.**
56
+
57
+ [![PyPI version](https://img.shields.io/pypi/v/schematize.svg)](https://pypi.org/project/schematize/)
58
+ [![Python](https://img.shields.io/pypi/pyversions/schematize.svg)](https://pypi.org/project/schematize/)
59
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
60
+ [![Docs](https://img.shields.io/badge/docs-pwr--ai.github.io-blue.svg)](https://pwr-ai.github.io/schematize)
61
+
62
+ [Documentation](https://pwr-ai.github.io/schematize) ·
63
+ [Quickstart](https://pwr-ai.github.io/schematize/quickstart/) ·
64
+ [Pipeline](https://pwr-ai.github.io/schematize/pipeline/) ·
65
+ [API reference](https://pwr-ai.github.io/schematize/api/schema_generator/)
66
+
67
+ </div>
68
+
69
+ ---
70
+
71
+ `schematize` is a Python library that turns a natural-language description of *what you want to
72
+ extract* into a **typed, validated extraction schema** — field names, types, descriptions, and
73
+ enums — ready to drive structured information extraction from a document collection.
74
+
75
+ Instead of hand-writing JSON schemas and discovering their gaps in production, you describe the
76
+ problem. A multi-agent [LangGraph](https://langchain-ai.github.io/langgraph/) pipeline asks
77
+ clarifying questions, drafts a schema, critiques and refines it, **tests it against real documents
78
+ from your corpus**, and opens a chat for final tweaks.
79
+
80
+ ## Why schematize?
81
+
82
+ - **Designing extraction schemas by hand is slow and brittle.** You guess the fields, miss edge
83
+ cases, and only find out when extraction quality is poor.
84
+ - **Raw "ask an LLM for a schema" gives you an untested first draft.** No critique loop, no contact
85
+ with your actual data, no typing guarantees.
86
+ - **schematize closes the loop:** clarify → draft → criteria-based refinement → **data-grounded
87
+ refinement against retrieved documents** → interactive chat. The output is a Pydantic model you
88
+ can plug straight into an extraction pipeline.
89
+ - **Use any LLM.** schematize accepts any LangChain chat model and any OpenAI-compatible endpoint, so
90
+ you can run it through a [LiteLLM](https://github.com/BerriAI/litellm) proxy — which is exactly how
91
+ we ran our experiments — and reach OpenAI, Anthropic, Gemini, or local models through one
92
+ interface. No provider lock-in.
93
+ - **Bring your own data.** Implementing a retriever is one async method; HuggingFace and Weaviate
94
+ adapters are built in. A schema-coverage evaluator is included too.
95
+
96
+ | | Hand-written schema | Ask-an-LLM once | **schematize** |
97
+ |---|:---:|:---:|:---:|
98
+ | Clarifies an ambiguous request | ❌ | ❌ | ✅ |
99
+ | Iterative critique & refinement | ❌ | ❌ | ✅ |
100
+ | Validated against **real documents** | ❌ | ❌ | ✅ |
101
+ | Typed Pydantic output | manual | ❌ | ✅ |
102
+ | Built-in coverage evaluation | ❌ | ❌ | ✅ |
103
+
104
+ ## Install
105
+
106
+ ```bash
107
+ pip install schematize
108
+ ```
109
+
110
+ Optional adapters and tooling:
111
+
112
+ ```bash
113
+ pip install "schematize[huggingface]" # FAISS retriever over HuggingFace datasets
114
+ pip install "schematize[weaviate]" # hybrid-search retriever for a Weaviate instance
115
+ pip install "schematize[scripts]" # Hydra-based CLI runners
116
+ ```
117
+
118
+ Requires Python 3.12+. Full options in the [installation guide](https://pwr-ai.github.io/schematize/installation/).
119
+
120
+ ## Quickstart
121
+
122
+ ```python
123
+ from langchain_openai import ChatOpenAI
124
+ from schematize import SchemaGenerator, load_prompts
125
+
126
+
127
+ # 1. Any object with this async __call__ is a valid retriever (DocumentRetriever protocol).
128
+ # Return documents relevant to `query`; each is shown to the data-assessment agent.
129
+ class MyRetriever:
130
+ async def __call__(self, query: str, max_docs: int = 100) -> list:
131
+ return [
132
+ {"text": "The court awarded 15,000 PLN in damages for breach of personal rights..."},
133
+ {"text": "Claim dismissed; the plaintiff failed to prove the violation..."},
134
+ ]
135
+
136
+
137
+ # 2. Load bundled prompts for your language/domain (en|pl × law|tax).
138
+ prompts = load_prompts(language="en", system_type="law")
139
+
140
+ llm = ChatOpenAI(model="gpt-4o", temperature=0.2)
141
+
142
+ generator = SchemaGenerator(llm=llm, retriever=MyRetriever(), **prompts)
143
+
144
+ # 3. Run the pipeline (interactive: it asks clarifying questions in the terminal).
145
+ state = generator.stream_graph_updates(
146
+ "Study personal-rights violations in civil cases and assess their severity."
147
+ )
148
+ print(state["current_schema"])
149
+ ```
150
+
151
+ A generated schema looks like this — a typed spec you can act on immediately:
152
+
153
+ ```python
154
+ {
155
+ "fields": [
156
+ {"name": "violation_type", "type_": "enum", "enum_name": "ViolationType",
157
+ "enum_values": ["privacy", "reputation", "image", "bodily_integrity"],
158
+ "description": "Category of personal right that was violated."},
159
+ {"name": "severity", "type_": "integer",
160
+ "description": "Severity of the violation on a 0–5 scale."},
161
+ {"name": "compensation_awarded", "type_": "boolean",
162
+ "description": "Whether monetary compensation was granted."},
163
+ {"name": "compensation_amount", "type_": "float",
164
+ "description": "Awarded amount in PLN, if any."},
165
+ ]
166
+ }
167
+ ```
168
+
169
+ Turn it into a Pydantic model and use it for extraction:
170
+
171
+ ```python
172
+ from schematize import SchemaFields, DynamicModelFactory
173
+
174
+ model_cls = DynamicModelFactory()(SchemaFields(**state["current_schema"]))
175
+ ```
176
+
177
+ ## Use any LLM (via LiteLLM)
178
+
179
+ `SchemaGenerator` takes any LangChain [`BaseChatModel`](https://python.langchain.com/docs/concepts/chat_models/).
180
+ Because it also honours an OpenAI-compatible `base_url`, the simplest way to reach **any** provider is
181
+ to put a [LiteLLM](https://github.com/BerriAI/litellm) proxy in front and point schematize at it —
182
+ the setup we used for our experiments:
183
+
184
+ ```python
185
+ from langchain_openai import ChatOpenAI
186
+
187
+ # Point at a LiteLLM proxy; the model name routes to OpenAI, Anthropic, Gemini, local, etc.
188
+ llm = ChatOpenAI(model="claude-opus-4-8", base_url="http://localhost:4000", api_key="sk-litellm")
189
+ ```
190
+
191
+ One interface, 100+ providers, no lock-in. See [Configuration](https://pwr-ai.github.io/schematize/configuration/).
192
+
193
+ ## Retrieval is pluggable
194
+
195
+ The core library has **no retrieval dependency**. Implementing your own retriever is a single async
196
+ method — the `DocumentRetriever` protocol shown in the quickstart — so you can wrap Elasticsearch,
197
+ Postgres FTS, a REST API, or any vector store.
198
+
199
+ Two adapters ship in the box:
200
+
201
+ **HuggingFace** (`[huggingface]`) — FAISS index over any HuggingFace dataset, cached to disk:
202
+
203
+ ```python
204
+ from schematize.retrieval.huggingface import HuggingFaceRetriever
205
+
206
+ retriever = HuggingFaceRetriever(
207
+ dataset_name="JuDDGES/pl-court-raw", text_column="text", index_path=".cache/court-index"
208
+ )
209
+ ```
210
+
211
+ **Weaviate** (`[weaviate]`) — hybrid search against a Weaviate instance:
212
+
213
+ ```python
214
+ from schematize.retrieval.weaviate import WeaviateRetriever
215
+
216
+ retriever = WeaviateRetriever(collection_name="LegalDocuments")
217
+ ```
218
+
219
+ Needs `WV_URL`, `WV_PORT`, `WV_GRPC_PORT`, `WV_API_KEY`. See the
220
+ [Weaviate guide](https://pwr-ai.github.io/schematize/guides/weaviate/).
221
+
222
+ `MMLWRobertaV2Retriever` (Polish-optimised, built on the HuggingFace base) is a worked **example of a
223
+ custom retriever** — read it as a template for specialising retrieval to your own model or language.
224
+ See the [custom retriever guide](https://pwr-ai.github.io/schematize/guides/custom-retriever/).
225
+
226
+ ## Evaluate a schema
227
+
228
+ Score how well a schema can answer a set of expert questions:
229
+
230
+ ```python
231
+ import yaml
232
+ from langchain_openai import ChatOpenAI
233
+ from schematize import SchemaEvaluator
234
+ from schematize.settings import PROMPTS_PATH
235
+
236
+ with open(PROMPTS_PATH / "eval" / "schema_evaluator.yaml") as f:
237
+ evaluation_prompt = yaml.safe_load(f)["schema_evaluator_prompt"]
238
+
239
+ evaluator = SchemaEvaluator(ChatOpenAI(model="gpt-4o"), evaluation_prompt)
240
+ result = evaluator.evaluate_schema(schema, questions=["How severe was the violation?", ...])
241
+ print(result.covered_questions, "/", result.total_questions)
242
+ ```
243
+
244
+ More in the [evaluation guide](https://pwr-ai.github.io/schematize/guides/evaluation/).
245
+
246
+ ## Command-line runners
247
+
248
+ With the `[scripts]` extra you get three console scripts:
249
+
250
+ ```bash
251
+ schematize-run # interactive pipeline
252
+ schematize-run-mocked +case=en_age # replay a stored case (no live prompts)
253
+ schematize-evaluate +case_name=age # evaluate against expert questions
254
+ ```
255
+
256
+ See the [CLI guide](https://pwr-ai.github.io/schematize/guides/cli/) for mocked-case files and options.
257
+
258
+ ## Reproducing our study
259
+
260
+ The experiments from our paper are driven by the mocked runner (schema generation) and the evaluator
261
+ (schema scoring against expert questions). Cases live in [`data/cases/`](data/cases/) and expert
262
+ question sets in [`data/eval/`](data/eval/) (`pl_age`, `pl_personal_rights`, `pl_medical_errors`).
263
+
264
+ ```bash
265
+ git clone https://github.com/pwr-ai/schematize && cd schematize
266
+ uv sync --extra scripts --extra huggingface
267
+
268
+ # Configure the LLM in a .env file (we used a LiteLLM proxy — see "Use any LLM" above)
269
+ printf 'API_KEY=...\nAPI_URL=...\n' > .env
270
+
271
+ # 1. Generate schemas for every case (multiple runs per case)
272
+ bash scripts/experiments/search_params.sh
273
+
274
+ # 2. Ablation over pipeline components (no problem-definition / no refinement / no data-grounding)
275
+ bash scripts/experiments/ablation.sh <model>
276
+
277
+ # 3. Evaluate generated schemas against expert questions
278
+ bash scripts/experiments/eval_multirun.sh <eval_model> <generation_model>
279
+ ```
280
+
281
+ The shell scripts in [`scripts/experiments/`](scripts/experiments/) are thin Hydra wrappers; edit the
282
+ `MODEL`/`CASES` variables at the top to change the grid. Results are written to the Hydra output
283
+ directory.
284
+
285
+ > **Note:** exact model names, seeds, and hyperparameters used in the paper are documented in the
286
+ > [reproduction guide](https://pwr-ai.github.io/schematize/) — fill in once the study is published.
287
+
288
+ ## Citation
289
+
290
+ If you use schematize in your research, please cite:
291
+
292
+ ```bibtex
293
+ TBA
294
+ ```
295
+
296
+ ## Development
297
+
298
+ ```bash
299
+ uv sync --extra dev
300
+ make check # ruff lint
301
+ make test # pytest + coverage
302
+ make fix # ruff --fix
303
+ ```
304
+
305
+ ## License
306
+
307
+ Released under the [MIT License](LICENSE).
@@ -0,0 +1,259 @@
1
+ <div align="center">
2
+
3
+ <img src="docs/assets/logo.png" alt="schematize logo" width="200" height="200">
4
+
5
+ # schematize
6
+
7
+ **Turn a plain-English research problem into a typed, data-tested extraction schema.**
8
+
9
+ [![PyPI version](https://img.shields.io/pypi/v/schematize.svg)](https://pypi.org/project/schematize/)
10
+ [![Python](https://img.shields.io/pypi/pyversions/schematize.svg)](https://pypi.org/project/schematize/)
11
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
12
+ [![Docs](https://img.shields.io/badge/docs-pwr--ai.github.io-blue.svg)](https://pwr-ai.github.io/schematize)
13
+
14
+ [Documentation](https://pwr-ai.github.io/schematize) ·
15
+ [Quickstart](https://pwr-ai.github.io/schematize/quickstart/) ·
16
+ [Pipeline](https://pwr-ai.github.io/schematize/pipeline/) ·
17
+ [API reference](https://pwr-ai.github.io/schematize/api/schema_generator/)
18
+
19
+ </div>
20
+
21
+ ---
22
+
23
+ `schematize` is a Python library that turns a natural-language description of *what you want to
24
+ extract* into a **typed, validated extraction schema** — field names, types, descriptions, and
25
+ enums — ready to drive structured information extraction from a document collection.
26
+
27
+ Instead of hand-writing JSON schemas and discovering their gaps in production, you describe the
28
+ problem. A multi-agent [LangGraph](https://langchain-ai.github.io/langgraph/) pipeline asks
29
+ clarifying questions, drafts a schema, critiques and refines it, **tests it against real documents
30
+ from your corpus**, and opens a chat for final tweaks.
31
+
32
+ ## Why schematize?
33
+
34
+ - **Designing extraction schemas by hand is slow and brittle.** You guess the fields, miss edge
35
+ cases, and only find out when extraction quality is poor.
36
+ - **Raw "ask an LLM for a schema" gives you an untested first draft.** No critique loop, no contact
37
+ with your actual data, no typing guarantees.
38
+ - **schematize closes the loop:** clarify → draft → criteria-based refinement → **data-grounded
39
+ refinement against retrieved documents** → interactive chat. The output is a Pydantic model you
40
+ can plug straight into an extraction pipeline.
41
+ - **Use any LLM.** schematize accepts any LangChain chat model and any OpenAI-compatible endpoint, so
42
+ you can run it through a [LiteLLM](https://github.com/BerriAI/litellm) proxy — which is exactly how
43
+ we ran our experiments — and reach OpenAI, Anthropic, Gemini, or local models through one
44
+ interface. No provider lock-in.
45
+ - **Bring your own data.** Implementing a retriever is one async method; HuggingFace and Weaviate
46
+ adapters are built in. A schema-coverage evaluator is included too.
47
+
48
+ | | Hand-written schema | Ask-an-LLM once | **schematize** |
49
+ |---|:---:|:---:|:---:|
50
+ | Clarifies an ambiguous request | ❌ | ❌ | ✅ |
51
+ | Iterative critique & refinement | ❌ | ❌ | ✅ |
52
+ | Validated against **real documents** | ❌ | ❌ | ✅ |
53
+ | Typed Pydantic output | manual | ❌ | ✅ |
54
+ | Built-in coverage evaluation | ❌ | ❌ | ✅ |
55
+
56
+ ## Install
57
+
58
+ ```bash
59
+ pip install schematize
60
+ ```
61
+
62
+ Optional adapters and tooling:
63
+
64
+ ```bash
65
+ pip install "schematize[huggingface]" # FAISS retriever over HuggingFace datasets
66
+ pip install "schematize[weaviate]" # hybrid-search retriever for a Weaviate instance
67
+ pip install "schematize[scripts]" # Hydra-based CLI runners
68
+ ```
69
+
70
+ Requires Python 3.12+. Full options in the [installation guide](https://pwr-ai.github.io/schematize/installation/).
71
+
72
+ ## Quickstart
73
+
74
+ ```python
75
+ from langchain_openai import ChatOpenAI
76
+ from schematize import SchemaGenerator, load_prompts
77
+
78
+
79
+ # 1. Any object with this async __call__ is a valid retriever (DocumentRetriever protocol).
80
+ # Return documents relevant to `query`; each is shown to the data-assessment agent.
81
+ class MyRetriever:
82
+ async def __call__(self, query: str, max_docs: int = 100) -> list:
83
+ return [
84
+ {"text": "The court awarded 15,000 PLN in damages for breach of personal rights..."},
85
+ {"text": "Claim dismissed; the plaintiff failed to prove the violation..."},
86
+ ]
87
+
88
+
89
+ # 2. Load bundled prompts for your language/domain (en|pl × law|tax).
90
+ prompts = load_prompts(language="en", system_type="law")
91
+
92
+ llm = ChatOpenAI(model="gpt-4o", temperature=0.2)
93
+
94
+ generator = SchemaGenerator(llm=llm, retriever=MyRetriever(), **prompts)
95
+
96
+ # 3. Run the pipeline (interactive: it asks clarifying questions in the terminal).
97
+ state = generator.stream_graph_updates(
98
+ "Study personal-rights violations in civil cases and assess their severity."
99
+ )
100
+ print(state["current_schema"])
101
+ ```
102
+
103
+ A generated schema looks like this — a typed spec you can act on immediately:
104
+
105
+ ```python
106
+ {
107
+ "fields": [
108
+ {"name": "violation_type", "type_": "enum", "enum_name": "ViolationType",
109
+ "enum_values": ["privacy", "reputation", "image", "bodily_integrity"],
110
+ "description": "Category of personal right that was violated."},
111
+ {"name": "severity", "type_": "integer",
112
+ "description": "Severity of the violation on a 0–5 scale."},
113
+ {"name": "compensation_awarded", "type_": "boolean",
114
+ "description": "Whether monetary compensation was granted."},
115
+ {"name": "compensation_amount", "type_": "float",
116
+ "description": "Awarded amount in PLN, if any."},
117
+ ]
118
+ }
119
+ ```
120
+
121
+ Turn it into a Pydantic model and use it for extraction:
122
+
123
+ ```python
124
+ from schematize import SchemaFields, DynamicModelFactory
125
+
126
+ model_cls = DynamicModelFactory()(SchemaFields(**state["current_schema"]))
127
+ ```
128
+
129
+ ## Use any LLM (via LiteLLM)
130
+
131
+ `SchemaGenerator` takes any LangChain [`BaseChatModel`](https://python.langchain.com/docs/concepts/chat_models/).
132
+ Because it also honours an OpenAI-compatible `base_url`, the simplest way to reach **any** provider is
133
+ to put a [LiteLLM](https://github.com/BerriAI/litellm) proxy in front and point schematize at it —
134
+ the setup we used for our experiments:
135
+
136
+ ```python
137
+ from langchain_openai import ChatOpenAI
138
+
139
+ # Point at a LiteLLM proxy; the model name routes to OpenAI, Anthropic, Gemini, local, etc.
140
+ llm = ChatOpenAI(model="claude-opus-4-8", base_url="http://localhost:4000", api_key="sk-litellm")
141
+ ```
142
+
143
+ One interface, 100+ providers, no lock-in. See [Configuration](https://pwr-ai.github.io/schematize/configuration/).
144
+
145
+ ## Retrieval is pluggable
146
+
147
+ The core library has **no retrieval dependency**. Implementing your own retriever is a single async
148
+ method — the `DocumentRetriever` protocol shown in the quickstart — so you can wrap Elasticsearch,
149
+ Postgres FTS, a REST API, or any vector store.
150
+
151
+ Two adapters ship in the box:
152
+
153
+ **HuggingFace** (`[huggingface]`) — FAISS index over any HuggingFace dataset, cached to disk:
154
+
155
+ ```python
156
+ from schematize.retrieval.huggingface import HuggingFaceRetriever
157
+
158
+ retriever = HuggingFaceRetriever(
159
+ dataset_name="JuDDGES/pl-court-raw", text_column="text", index_path=".cache/court-index"
160
+ )
161
+ ```
162
+
163
+ **Weaviate** (`[weaviate]`) — hybrid search against a Weaviate instance:
164
+
165
+ ```python
166
+ from schematize.retrieval.weaviate import WeaviateRetriever
167
+
168
+ retriever = WeaviateRetriever(collection_name="LegalDocuments")
169
+ ```
170
+
171
+ Needs `WV_URL`, `WV_PORT`, `WV_GRPC_PORT`, `WV_API_KEY`. See the
172
+ [Weaviate guide](https://pwr-ai.github.io/schematize/guides/weaviate/).
173
+
174
+ `MMLWRobertaV2Retriever` (Polish-optimised, built on the HuggingFace base) is a worked **example of a
175
+ custom retriever** — read it as a template for specialising retrieval to your own model or language.
176
+ See the [custom retriever guide](https://pwr-ai.github.io/schematize/guides/custom-retriever/).
177
+
178
+ ## Evaluate a schema
179
+
180
+ Score how well a schema can answer a set of expert questions:
181
+
182
+ ```python
183
+ import yaml
184
+ from langchain_openai import ChatOpenAI
185
+ from schematize import SchemaEvaluator
186
+ from schematize.settings import PROMPTS_PATH
187
+
188
+ with open(PROMPTS_PATH / "eval" / "schema_evaluator.yaml") as f:
189
+ evaluation_prompt = yaml.safe_load(f)["schema_evaluator_prompt"]
190
+
191
+ evaluator = SchemaEvaluator(ChatOpenAI(model="gpt-4o"), evaluation_prompt)
192
+ result = evaluator.evaluate_schema(schema, questions=["How severe was the violation?", ...])
193
+ print(result.covered_questions, "/", result.total_questions)
194
+ ```
195
+
196
+ More in the [evaluation guide](https://pwr-ai.github.io/schematize/guides/evaluation/).
197
+
198
+ ## Command-line runners
199
+
200
+ With the `[scripts]` extra you get three console scripts:
201
+
202
+ ```bash
203
+ schematize-run # interactive pipeline
204
+ schematize-run-mocked +case=en_age # replay a stored case (no live prompts)
205
+ schematize-evaluate +case_name=age # evaluate against expert questions
206
+ ```
207
+
208
+ See the [CLI guide](https://pwr-ai.github.io/schematize/guides/cli/) for mocked-case files and options.
209
+
210
+ ## Reproducing our study
211
+
212
+ The experiments from our paper are driven by the mocked runner (schema generation) and the evaluator
213
+ (schema scoring against expert questions). Cases live in [`data/cases/`](data/cases/) and expert
214
+ question sets in [`data/eval/`](data/eval/) (`pl_age`, `pl_personal_rights`, `pl_medical_errors`).
215
+
216
+ ```bash
217
+ git clone https://github.com/pwr-ai/schematize && cd schematize
218
+ uv sync --extra scripts --extra huggingface
219
+
220
+ # Configure the LLM in a .env file (we used a LiteLLM proxy — see "Use any LLM" above)
221
+ printf 'API_KEY=...\nAPI_URL=...\n' > .env
222
+
223
+ # 1. Generate schemas for every case (multiple runs per case)
224
+ bash scripts/experiments/search_params.sh
225
+
226
+ # 2. Ablation over pipeline components (no problem-definition / no refinement / no data-grounding)
227
+ bash scripts/experiments/ablation.sh <model>
228
+
229
+ # 3. Evaluate generated schemas against expert questions
230
+ bash scripts/experiments/eval_multirun.sh <eval_model> <generation_model>
231
+ ```
232
+
233
+ The shell scripts in [`scripts/experiments/`](scripts/experiments/) are thin Hydra wrappers; edit the
234
+ `MODEL`/`CASES` variables at the top to change the grid. Results are written to the Hydra output
235
+ directory.
236
+
237
+ > **Note:** exact model names, seeds, and hyperparameters used in the paper are documented in the
238
+ > [reproduction guide](https://pwr-ai.github.io/schematize/) — fill in once the study is published.
239
+
240
+ ## Citation
241
+
242
+ If you use schematize in your research, please cite:
243
+
244
+ ```bibtex
245
+ TBA
246
+ ```
247
+
248
+ ## Development
249
+
250
+ ```bash
251
+ uv sync --extra dev
252
+ make check # ruff lint
253
+ make test # pytest + coverage
254
+ make fix # ruff --fix
255
+ ```
256
+
257
+ ## License
258
+
259
+ Released under the [MIT License](LICENSE).