kroma-om 0.2.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 (57) hide show
  1. kroma_om-0.2.0/LICENSE +28 -0
  2. kroma_om-0.2.0/PKG-INFO +165 -0
  3. kroma_om-0.2.0/README.md +118 -0
  4. kroma_om-0.2.0/pyproject.toml +70 -0
  5. kroma_om-0.2.0/setup.cfg +4 -0
  6. kroma_om-0.2.0/src/kroma/__init__.py +36 -0
  7. kroma_om-0.2.0/src/kroma/agents/__init__.py +1 -0
  8. kroma_om-0.2.0/src/kroma/agents/factory.py +43 -0
  9. kroma_om-0.2.0/src/kroma/agents/inference.py +211 -0
  10. kroma_om-0.2.0/src/kroma/agents/utils.py +80 -0
  11. kroma_om-0.2.0/src/kroma/algorithms/__init__.py +1 -0
  12. kroma_om-0.2.0/src/kroma/algorithms/node2vec.py +135 -0
  13. kroma_om-0.2.0/src/kroma/algorithms/refinement.py +192 -0
  14. kroma_om-0.2.0/src/kroma/algorithms/utils.py +55 -0
  15. kroma_om-0.2.0/src/kroma/alignment.py +128 -0
  16. kroma_om-0.2.0/src/kroma/cli.py +184 -0
  17. kroma_om-0.2.0/src/kroma/config/__init__.py +1 -0
  18. kroma_om-0.2.0/src/kroma/config/arguments.py +54 -0
  19. kroma_om-0.2.0/src/kroma/config/constants.py +136 -0
  20. kroma_om-0.2.0/src/kroma/costs.py +65 -0
  21. kroma_om-0.2.0/src/kroma/decorators.py +84 -0
  22. kroma_om-0.2.0/src/kroma/inference/__init__.py +1 -0
  23. kroma_om-0.2.0/src/kroma/inference/factory.py +148 -0
  24. kroma_om-0.2.0/src/kroma/matcher.py +363 -0
  25. kroma_om-0.2.0/src/kroma/metrics/__init__.py +1 -0
  26. kroma_om-0.2.0/src/kroma/metrics/scoring.py +41 -0
  27. kroma_om-0.2.0/src/kroma/model_providers/__init__.py +1 -0
  28. kroma_om-0.2.0/src/kroma/model_providers/base_model.py +176 -0
  29. kroma_om-0.2.0/src/kroma/model_providers/chat_providers.py +101 -0
  30. kroma_om-0.2.0/src/kroma/model_providers/embedding_providers.py +62 -0
  31. kroma_om-0.2.0/src/kroma/model_providers/providers.py +61 -0
  32. kroma_om-0.2.0/src/kroma/observers.py +46 -0
  33. kroma_om-0.2.0/src/kroma/ontology.py +195 -0
  34. kroma_om-0.2.0/src/kroma/pipeline.py +394 -0
  35. kroma_om-0.2.0/src/kroma/repository.py +88 -0
  36. kroma_om-0.2.0/src/kroma/retrieval/__init__.py +1 -0
  37. kroma_om-0.2.0/src/kroma/retrieval/vector_store.py +61 -0
  38. kroma_om-0.2.0/src/kroma/tasks/__init__.py +1 -0
  39. kroma_om-0.2.0/src/kroma/tasks/builder.py +122 -0
  40. kroma_om-0.2.0/src/kroma/tasks/loader.py +92 -0
  41. kroma_om-0.2.0/src/kroma/tasks/prompts/__init__.py +1 -0
  42. kroma_om-0.2.0/src/kroma/tasks/prompts/builders.py +174 -0
  43. kroma_om-0.2.0/src/kroma/tasks/prompts/templates.py +44 -0
  44. kroma_om-0.2.0/src/kroma/tasks/query_engine.py +314 -0
  45. kroma_om-0.2.0/src/kroma/tasks/query_templates.py +223 -0
  46. kroma_om-0.2.0/src/kroma/utils/__init__.py +1 -0
  47. kroma_om-0.2.0/src/kroma/utils/env.py +25 -0
  48. kroma_om-0.2.0/src/kroma/utils/file.py +141 -0
  49. kroma_om-0.2.0/src/kroma_om.egg-info/PKG-INFO +165 -0
  50. kroma_om-0.2.0/src/kroma_om.egg-info/SOURCES.txt +55 -0
  51. kroma_om-0.2.0/src/kroma_om.egg-info/dependency_links.txt +1 -0
  52. kroma_om-0.2.0/src/kroma_om.egg-info/entry_points.txt +2 -0
  53. kroma_om-0.2.0/src/kroma_om.egg-info/requires.txt +25 -0
  54. kroma_om-0.2.0/src/kroma_om.egg-info/top_level.txt +1 -0
  55. kroma_om-0.2.0/tests/test_api.py +187 -0
  56. kroma_om-0.2.0/tests/test_costs.py +13 -0
  57. kroma_om-0.2.0/tests/test_matching.py +63 -0
kroma_om-0.2.0/LICENSE ADDED
@@ -0,0 +1,28 @@
1
+ BSD 3-Clause License
2
+
3
+ Copyright (c) 2025, Lam Nguyen, Erika Barcelos, Roger French, and Yinghui Wu
4
+
5
+ Redistribution and use in source and binary forms, with or without
6
+ modification, are permitted provided that the following conditions are met:
7
+
8
+ 1. Redistributions of source code must retain the above copyright notice, this
9
+ list of conditions and the following disclaimer.
10
+
11
+ 2. Redistributions in binary form must reproduce the above copyright notice,
12
+ this list of conditions and the following disclaimer in the documentation
13
+ and/or other materials provided with the distribution.
14
+
15
+ 3. Neither the name of the copyright holder nor the names of its
16
+ contributors may be used to endorse or promote products derived from
17
+ this software without specific prior written permission.
18
+
19
+ THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
20
+ AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
21
+ IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
22
+ DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
23
+ FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
24
+ DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
25
+ SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
26
+ CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
27
+ OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
28
+ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
@@ -0,0 +1,165 @@
1
+ Metadata-Version: 2.4
2
+ Name: kroma-om
3
+ Version: 0.2.0
4
+ Summary: Ontology matching with knowledge retrieval and large language models
5
+ Author: Lam Nguyen, Erika Barcelos, Roger French, Yinghui Wu
6
+ License-Expression: BSD-3-Clause
7
+ Project-URL: Homepage, https://github.com/cwru-sdle/kroma
8
+ Project-URL: Documentation, https://kroma.readthedocs.io/en/latest/
9
+ Project-URL: Repository, https://github.com/cwru-sdle/kroma
10
+ Project-URL: Issues, https://github.com/cwru-sdle/kroma/issues
11
+ Project-URL: Paper, https://doi.org/10.1007/978-3-032-09527-5_34
12
+ Project-URL: arXiv, https://arxiv.org/abs/2507.14032
13
+ Keywords: ontology matching,ontology alignment,knowledge graphs,large language models,retrieval-augmented generation,semantic web
14
+ Classifier: Development Status :: 4 - Beta
15
+ Classifier: Intended Audience :: Science/Research
16
+ Classifier: Operating System :: OS Independent
17
+ Classifier: Programming Language :: Python :: 3
18
+ Classifier: Programming Language :: Python :: 3.12
19
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
20
+ Classifier: Topic :: Scientific/Engineering :: Information Analysis
21
+ Requires-Python: >=3.12
22
+ Description-Content-Type: text/markdown
23
+ License-File: LICENSE
24
+ Requires-Dist: numpy
25
+ Requires-Dist: python-dotenv
26
+ Requires-Dist: rdflib
27
+ Requires-Dist: requests
28
+ Requires-Dist: scipy
29
+ Provides-Extra: models
30
+ Requires-Dist: nltk; extra == "models"
31
+ Requires-Dist: networkx; extra == "models"
32
+ Requires-Dist: node2vec; extra == "models"
33
+ Requires-Dist: openai; extra == "models"
34
+ Requires-Dist: scikit-learn; extra == "models"
35
+ Requires-Dist: scipy; extra == "models"
36
+ Requires-Dist: sentence-transformers; extra == "models"
37
+ Requires-Dist: SPARQLWrapper; extra == "models"
38
+ Requires-Dist: together; extra == "models"
39
+ Requires-Dist: torch; extra == "models"
40
+ Requires-Dist: transformers; extra == "models"
41
+ Provides-Extra: dev
42
+ Requires-Dist: pytest; extra == "dev"
43
+ Provides-Extra: docs
44
+ Requires-Dist: sphinx>=7; extra == "docs"
45
+ Requires-Dist: sphinx-rtd-theme>=2; extra == "docs"
46
+ Dynamic: license-file
47
+
48
+ <div align="center">
49
+
50
+ # KROMA
51
+
52
+ **Ontology matching with knowledge retrieval and large language models**
53
+
54
+ [![PyPI](https://img.shields.io/pypi/v/kroma-om)](https://pypi.org/project/kroma-om/)
55
+ [![Docs](https://readthedocs.org/projects/kroma/badge/?version=latest)](https://kroma.readthedocs.io/en/latest/)
56
+ [![ISWC 2025](https://img.shields.io/badge/ISWC-2025-b31b1b)](https://doi.org/10.1007/978-3-032-09527-5_34)
57
+ [![License: BSD 3-Clause](https://img.shields.io/badge/license-BSD%203--Clause-green.svg)](https://github.com/cwru-sdle/kroma/blob/main/LICENSE)
58
+
59
+ [**Documentation**](https://kroma.readthedocs.io/en/latest/) ·
60
+ [Quick start](#quick-start) ·
61
+ [Example](#example) ·
62
+ [Results](#results) ·
63
+ [Paper](https://doi.org/10.1007/978-3-032-09527-5_34) ·
64
+ [Citation](#citation)
65
+
66
+ </div>
67
+
68
+ KROMA finds the concepts that two ontologies share. It retrieves likely
69
+ matches for each concept, accepts identical names directly, and asks a large
70
+ language model about the rest, showing it each concept's labels, synonyms,
71
+ parents and children.
72
+
73
+ ## Quick start
74
+
75
+ ```bash
76
+ pip install kroma-om openai
77
+ export OPENAI_API_KEY=sk-...
78
+
79
+ kroma match mouse.owl human.owl -o alignments.csv
80
+ ```
81
+
82
+ Or from Python:
83
+
84
+ ```python
85
+ import kroma
86
+
87
+ result = kroma.match("mouse.owl", "human.owl", model="gpt-4o-mini")
88
+ result.save("alignments.csv") # .csv, .tsv, or .rdf (OAEI format)
89
+ ```
90
+
91
+ KROMA reads any RDF format (`.owl`, `.ttl`, `.rdf`, …) and works with OpenAI,
92
+ Together AI, or any function that takes a prompt and returns text. See the
93
+ [documentation](https://kroma.readthedocs.io/en/latest/) for all
94
+ options.
95
+
96
+ ## Example
97
+
98
+ Try KROMA with a local [Ollama](https://ollama.com) model, no API key needed:
99
+
100
+ ```bash
101
+ git clone https://github.com/cwru-sdle/kroma.git && cd kroma/examples
102
+ pip install kroma-om
103
+ ollama pull llama3.2:1b
104
+ python match_ollama.py --model llama3.2:1b
105
+ ```
106
+
107
+ | Source | Target | Confidence (0-10) |
108
+ |:--|:--|:--:|
109
+ | heart | cardiac organ | 10 |
110
+ | lung | lungs | 8 |
111
+ | organ | body organ | 8 |
112
+
113
+ Run with `llama3.2:1b` on a small demo dataset; see the
114
+ [example](https://github.com/cwru-sdle/kroma/tree/main/examples) for details.
115
+
116
+ ## Results
117
+
118
+ F1 (%) on six OAEI tracks, from the [ISWC 2025 paper](https://doi.org/10.1007/978-3-032-09527-5_34):
119
+
120
+ | Track | KROMA | Best baseline |
121
+ |:--|:--:|:--:|
122
+ | Mouse–Human | **94.94** | 92.20 |
123
+ | NCIT–DOID | **98.63** | 94.80 |
124
+ | Nell–DBpedia | **97.08** | 96.00 |
125
+ | YAGO–Wikidata | **95.54** | 92.19 |
126
+ | ENVO–SWEET | **93.18** | 83.70 |
127
+ | MI–MatOnto | **61.25** | 32.97 |
128
+
129
+ Ablations, other models and cost are in the
130
+ [full paper](https://github.com/cwru-sdle/kroma/blob/main/full.pdf).
131
+
132
+ ## Reproducing the paper
133
+
134
+ ```bash
135
+ git clone https://github.com/cwru-sdle/kroma.git && cd kroma
136
+ ./setup.sh && source .venv/bin/activate
137
+ uv pip install ".[models]"
138
+
139
+ kroma run --method_config kroma_scibert_envo_sweet \
140
+ --llm meta-llama/Llama-3.3-70B-Instruct-Turbo-Free
141
+ ```
142
+
143
+ Datasets, configs and flags are described in
144
+ [Benchmark experiments](https://kroma.readthedocs.io/en/latest/usage.html).
145
+
146
+ ## Citation
147
+
148
+ ```bibtex
149
+ @inproceedings{nguyen_2025_kroma,
150
+ title = {KROMA: Ontology Matching with Knowledge Retrieval and Large Language Models},
151
+ author = {Lam Nguyen and Erika Barcelos and Roger French and Yinghui Wu},
152
+ booktitle = {The Semantic Web -- ISWC 2025: 24th International Semantic Web Conference, Proceedings, Part I},
153
+ series = {Lecture Notes in Computer Science},
154
+ publisher = {Springer},
155
+ pages = {629--649},
156
+ year = {2025},
157
+ doi = {10.1007/978-3-032-09527-5_34}
158
+ }
159
+ ```
160
+
161
+ ## License
162
+
163
+ [BSD 3-Clause](https://github.com/cwru-sdle/kroma/blob/main/LICENSE) ·
164
+ Developed at the [SDLE Research Center](https://github.com/cwru-sdle),
165
+ Case Western Reserve University.
@@ -0,0 +1,118 @@
1
+ <div align="center">
2
+
3
+ # KROMA
4
+
5
+ **Ontology matching with knowledge retrieval and large language models**
6
+
7
+ [![PyPI](https://img.shields.io/pypi/v/kroma-om)](https://pypi.org/project/kroma-om/)
8
+ [![Docs](https://readthedocs.org/projects/kroma/badge/?version=latest)](https://kroma.readthedocs.io/en/latest/)
9
+ [![ISWC 2025](https://img.shields.io/badge/ISWC-2025-b31b1b)](https://doi.org/10.1007/978-3-032-09527-5_34)
10
+ [![License: BSD 3-Clause](https://img.shields.io/badge/license-BSD%203--Clause-green.svg)](https://github.com/cwru-sdle/kroma/blob/main/LICENSE)
11
+
12
+ [**Documentation**](https://kroma.readthedocs.io/en/latest/) ·
13
+ [Quick start](#quick-start) ·
14
+ [Example](#example) ·
15
+ [Results](#results) ·
16
+ [Paper](https://doi.org/10.1007/978-3-032-09527-5_34) ·
17
+ [Citation](#citation)
18
+
19
+ </div>
20
+
21
+ KROMA finds the concepts that two ontologies share. It retrieves likely
22
+ matches for each concept, accepts identical names directly, and asks a large
23
+ language model about the rest, showing it each concept's labels, synonyms,
24
+ parents and children.
25
+
26
+ ## Quick start
27
+
28
+ ```bash
29
+ pip install kroma-om openai
30
+ export OPENAI_API_KEY=sk-...
31
+
32
+ kroma match mouse.owl human.owl -o alignments.csv
33
+ ```
34
+
35
+ Or from Python:
36
+
37
+ ```python
38
+ import kroma
39
+
40
+ result = kroma.match("mouse.owl", "human.owl", model="gpt-4o-mini")
41
+ result.save("alignments.csv") # .csv, .tsv, or .rdf (OAEI format)
42
+ ```
43
+
44
+ KROMA reads any RDF format (`.owl`, `.ttl`, `.rdf`, …) and works with OpenAI,
45
+ Together AI, or any function that takes a prompt and returns text. See the
46
+ [documentation](https://kroma.readthedocs.io/en/latest/) for all
47
+ options.
48
+
49
+ ## Example
50
+
51
+ Try KROMA with a local [Ollama](https://ollama.com) model, no API key needed:
52
+
53
+ ```bash
54
+ git clone https://github.com/cwru-sdle/kroma.git && cd kroma/examples
55
+ pip install kroma-om
56
+ ollama pull llama3.2:1b
57
+ python match_ollama.py --model llama3.2:1b
58
+ ```
59
+
60
+ | Source | Target | Confidence (0-10) |
61
+ |:--|:--|:--:|
62
+ | heart | cardiac organ | 10 |
63
+ | lung | lungs | 8 |
64
+ | organ | body organ | 8 |
65
+
66
+ Run with `llama3.2:1b` on a small demo dataset; see the
67
+ [example](https://github.com/cwru-sdle/kroma/tree/main/examples) for details.
68
+
69
+ ## Results
70
+
71
+ F1 (%) on six OAEI tracks, from the [ISWC 2025 paper](https://doi.org/10.1007/978-3-032-09527-5_34):
72
+
73
+ | Track | KROMA | Best baseline |
74
+ |:--|:--:|:--:|
75
+ | Mouse–Human | **94.94** | 92.20 |
76
+ | NCIT–DOID | **98.63** | 94.80 |
77
+ | Nell–DBpedia | **97.08** | 96.00 |
78
+ | YAGO–Wikidata | **95.54** | 92.19 |
79
+ | ENVO–SWEET | **93.18** | 83.70 |
80
+ | MI–MatOnto | **61.25** | 32.97 |
81
+
82
+ Ablations, other models and cost are in the
83
+ [full paper](https://github.com/cwru-sdle/kroma/blob/main/full.pdf).
84
+
85
+ ## Reproducing the paper
86
+
87
+ ```bash
88
+ git clone https://github.com/cwru-sdle/kroma.git && cd kroma
89
+ ./setup.sh && source .venv/bin/activate
90
+ uv pip install ".[models]"
91
+
92
+ kroma run --method_config kroma_scibert_envo_sweet \
93
+ --llm meta-llama/Llama-3.3-70B-Instruct-Turbo-Free
94
+ ```
95
+
96
+ Datasets, configs and flags are described in
97
+ [Benchmark experiments](https://kroma.readthedocs.io/en/latest/usage.html).
98
+
99
+ ## Citation
100
+
101
+ ```bibtex
102
+ @inproceedings{nguyen_2025_kroma,
103
+ title = {KROMA: Ontology Matching with Knowledge Retrieval and Large Language Models},
104
+ author = {Lam Nguyen and Erika Barcelos and Roger French and Yinghui Wu},
105
+ booktitle = {The Semantic Web -- ISWC 2025: 24th International Semantic Web Conference, Proceedings, Part I},
106
+ series = {Lecture Notes in Computer Science},
107
+ publisher = {Springer},
108
+ pages = {629--649},
109
+ year = {2025},
110
+ doi = {10.1007/978-3-032-09527-5_34}
111
+ }
112
+ ```
113
+
114
+ ## License
115
+
116
+ [BSD 3-Clause](https://github.com/cwru-sdle/kroma/blob/main/LICENSE) ·
117
+ Developed at the [SDLE Research Center](https://github.com/cwru-sdle),
118
+ Case Western Reserve University.
@@ -0,0 +1,70 @@
1
+ [build-system]
2
+ requires = ["setuptools>=77.0.3"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "kroma-om"
7
+ version = "0.2.0"
8
+ description = "Ontology matching with knowledge retrieval and large language models"
9
+ readme = "README.md"
10
+ requires-python = ">=3.12"
11
+ authors = [
12
+ { name = "Lam Nguyen" },
13
+ { name = "Erika Barcelos" },
14
+ { name = "Roger French" },
15
+ { name = "Yinghui Wu" },
16
+ ]
17
+ license = "BSD-3-Clause"
18
+ keywords = ["ontology matching", "ontology alignment", "knowledge graphs", "large language models", "retrieval-augmented generation", "semantic web"]
19
+ classifiers = [
20
+ "Development Status :: 4 - Beta",
21
+ "Intended Audience :: Science/Research",
22
+ "Operating System :: OS Independent",
23
+ "Programming Language :: Python :: 3",
24
+ "Programming Language :: Python :: 3.12",
25
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
26
+ "Topic :: Scientific/Engineering :: Information Analysis",
27
+ ]
28
+ dependencies = [
29
+ "numpy",
30
+ "python-dotenv",
31
+ "rdflib",
32
+ "requests",
33
+ "scipy",
34
+ ]
35
+
36
+ [project.optional-dependencies]
37
+ models = [
38
+ "nltk",
39
+ "networkx",
40
+ "node2vec",
41
+ "openai",
42
+ "scikit-learn",
43
+ "scipy",
44
+ "sentence-transformers",
45
+ "SPARQLWrapper",
46
+ "together",
47
+ "torch",
48
+ "transformers",
49
+ ]
50
+ dev = ["pytest"]
51
+ docs = ["sphinx>=7", "sphinx-rtd-theme>=2"]
52
+
53
+ [project.scripts]
54
+ kroma = "kroma.cli:main"
55
+
56
+ [project.urls]
57
+ Homepage = "https://github.com/cwru-sdle/kroma"
58
+ Documentation = "https://kroma.readthedocs.io/en/latest/"
59
+ Repository = "https://github.com/cwru-sdle/kroma"
60
+ Issues = "https://github.com/cwru-sdle/kroma/issues"
61
+ Paper = "https://doi.org/10.1007/978-3-032-09527-5_34"
62
+ arXiv = "https://arxiv.org/abs/2507.14032"
63
+
64
+ [tool.setuptools.packages.find]
65
+ where = ["src"]
66
+ include = ["kroma*"]
67
+
68
+ [tool.pytest.ini_options]
69
+ testpaths = ["tests"]
70
+ pythonpath = ["src"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,36 @@
1
+ """KROMA: ontology matching with knowledge retrieval and large language models.
2
+
3
+ KROMA decides whether a concept in a source ontology and a concept in a target
4
+ ontology refer to the same real-world entity. For each candidate pair it
5
+ retrieves the concepts' context (ontology structure, cached dictionaries,
6
+ SPARQL endpoints and ConceptNet), asks one or more language models for a
7
+ Yes/No answer with a confidence, and refines a bisimulation-based partition of
8
+ the merged ontology graph as answers arrive.
9
+
10
+ To align two ontologies, call :func:`kroma.match` (or run ``kroma match``).
11
+ To reproduce the paper's benchmarks, use ``kroma run`` and
12
+ :class:`kroma.pipeline.MatchingRun`.
13
+
14
+ Example:
15
+ >>> import kroma
16
+ >>> result = kroma.match("source.owl", "target.owl") # doctest: +SKIP
17
+ >>> result.save("alignments.csv") # doctest: +SKIP
18
+ """
19
+
20
+ __version__ = "0.2.0"
21
+
22
+ from kroma.alignment import Alignment, MatchResult, save_alignments
23
+ from kroma.matcher import find_candidates, match
24
+ from kroma.ontology import Concept, Ontology, load_ontology
25
+
26
+ __all__ = [
27
+ "__version__",
28
+ "match",
29
+ "find_candidates",
30
+ "load_ontology",
31
+ "Ontology",
32
+ "Concept",
33
+ "Alignment",
34
+ "MatchResult",
35
+ "save_alignments",
36
+ ]
@@ -0,0 +1 @@
1
+ """Language-model agents: creation, prompting and answer parsing."""
@@ -0,0 +1,43 @@
1
+ """Create the chat agents used for single-model and debate inference."""
2
+
3
+ from typing import Dict, Any, List
4
+ from kroma.inference.factory import create_inference_model
5
+
6
+ def create_agents_from_config(
7
+ configs: List[Dict[str, str]]
8
+ ) -> Dict[int, Any]:
9
+ """Create one agent per config entry.
10
+
11
+ Args:
12
+ configs (list[dict[str, str]]): Entries with ``agent_type`` (a chat
13
+ backend) and ``agent_name`` (a model), for example
14
+ ``[{"agent_type": "openai", "agent_name": "gpt-4o-mini"}]``.
15
+
16
+ Returns:
17
+ dict[int, BaseModel]: Agents keyed by their position in ``configs``.
18
+ """
19
+ agents: Dict[int, Any] = {}
20
+ for idx, cfg in enumerate(configs):
21
+ agents[idx] = create_inference_model(cfg['agent_type'], cfg['agent_name'])
22
+ return agents
23
+
24
+
25
+ def create_agents(
26
+ n_agents: int,
27
+ backend: str,
28
+ model_name: str
29
+ ) -> Dict[int, Any]:
30
+ """Create ``n_agents`` agents with the same backend and model.
31
+
32
+ Args:
33
+ n_agents (int): Number of agents.
34
+ backend (str): Chat backend name, such as ``"openai"``.
35
+ model_name (str): Model name.
36
+
37
+ Returns:
38
+ dict[int, BaseModel]: Agents keyed ``0`` to ``n_agents - 1``.
39
+ """
40
+ return {
41
+ i: create_inference_model(backend, model_name)
42
+ for i in range(n_agents)
43
+ }
@@ -0,0 +1,211 @@
1
+ """Single-model and debate inference strategies."""
2
+
3
+ from abc import ABC, abstractmethod
4
+ from typing import Any, Dict, List, Tuple
5
+
6
+ from kroma.agents.factory import create_agents, create_agents_from_config
7
+ from kroma.agents.utils import (
8
+ active_learning_score,
9
+ extract_confidence,
10
+ extract_yes_no,
11
+ majority_vote,
12
+ )
13
+ from kroma.tasks.prompts.builders import (
14
+ build_last_round_prompt,
15
+ build_round_prompt,
16
+ build_task_prompt,
17
+ )
18
+ from kroma.tasks.query_engine import query
19
+
20
+
21
+ def expand_term(term, pcmap, dictionary, options):
22
+ """Gather a concept's context as a prompt-ready tuple.
23
+
24
+ Args:
25
+ term (tuple): ``(code, ...)``; only the code is used.
26
+ pcmap (dict): Parent and child maps for the concept's ontology.
27
+ dictionary (dict): Cached dictionaries.
28
+ options (list[str]): Context sources for
29
+ :func:`kroma.tasks.query_engine.query`.
30
+
31
+ Returns:
32
+ tuple: ``(code, parents, children, synonyms, labels)``.
33
+ """
34
+ parents, children, synonyms, labels = query(term[0], pcmap, dictionary, options)
35
+ return term[0], parents, children, synonyms, labels
36
+
37
+
38
+ def _run_rounds(agents, n_rounds, src, tgt, context_block, reasoning):
39
+ history: List[Tuple[int, int]] = []
40
+ metrics = {"input_token": 0, "output_token": 0, "api_calls": 0}
41
+ for round_idx in range(n_rounds):
42
+ if round_idx == 0:
43
+ system_prompt, base_prompt = build_task_prompt(src, tgt, reasoning=reasoning)
44
+ elif round_idx == n_rounds - 1:
45
+ system_prompt, base_prompt = build_last_round_prompt(src, tgt, history, perturb=False)
46
+ else:
47
+ system_prompt, base_prompt = build_round_prompt(src, tgt, history, perturb=False)
48
+ user_prompt = context_block + base_prompt
49
+ for agent in agents.values():
50
+ response, input_tokens, output_tokens = agent.generate(
51
+ system_prompt,
52
+ user_prompt,
53
+ include_tokens=True,
54
+ truncate=False,
55
+ )
56
+ history.append((extract_yes_no(response), extract_confidence(response)))
57
+ metrics["input_token"] += input_tokens
58
+ metrics["output_token"] += output_tokens
59
+ metrics["api_calls"] += 1
60
+ prediction, confidence = majority_vote(history)
61
+ return prediction, metrics, confidence
62
+
63
+
64
+ class InferenceStrategy(ABC):
65
+ """Chooses agents and how many debate rounds to run."""
66
+
67
+ @abstractmethod
68
+ def agents(self, backend, model_name, agent_configs):
69
+ """Return the agent mapping and the number of rounds.
70
+
71
+ Returns:
72
+ tuple[dict[int, BaseModel], int | None]: Agents and rounds;
73
+ ``None`` rounds means use the ``n_rounds`` passed to :meth:`run`.
74
+ """
75
+
76
+ def run(
77
+ self,
78
+ backend,
79
+ model_name,
80
+ source_term,
81
+ target_term,
82
+ pcmaps,
83
+ dictionary,
84
+ options,
85
+ agent_configs,
86
+ n_rounds=3,
87
+ reasoning=False,
88
+ active_learning=False,
89
+ f1_score=0.0,
90
+ context=None,
91
+ ):
92
+ """Score one source/target pair.
93
+
94
+ Every agent answers in every round; the prediction is the majority
95
+ answer and the confidence its average (see
96
+ :func:`kroma.agents.utils.majority_vote`).
97
+
98
+ Args:
99
+ backend (str): Chat backend name.
100
+ model_name (str): Model name.
101
+ source_term (tuple): Source ``(code, ...)``.
102
+ target_term (tuple): Target ``(code, ...)``.
103
+ pcmaps (dict): ``{"source": pcmap, "target": pcmap}``.
104
+ dictionary (dict): Cached dictionaries.
105
+ options (list[str]): Context sources.
106
+ agent_configs (list[dict]): Agent configs for debate.
107
+ n_rounds (int): Rounds when the strategy does not fix them.
108
+ reasoning (bool): Passed to the prompt builder.
109
+ active_learning (bool): Accept the answer only when
110
+ :func:`~kroma.agents.utils.active_learning_score` is positive.
111
+ f1_score (float): Running F1, which scales the active-learning
112
+ noise.
113
+ context (str | None): Text placed before every user prompt.
114
+
115
+ Returns:
116
+ tuple: ``(prediction, metrics, confidence, accept)``. ``prediction``
117
+ is 1 or 0, ``metrics`` holds ``input_token``, ``output_token`` and
118
+ ``api_calls``, ``confidence`` is 0 to 10, and ``accept`` says
119
+ whether to refine the graph with this answer.
120
+ """
121
+ agents, rounds = self.agents(backend, model_name, agent_configs)
122
+ if rounds is None:
123
+ rounds = n_rounds
124
+ source = expand_term(source_term, pcmaps["source"], dictionary, options)
125
+ target = expand_term(target_term, pcmaps["target"], dictionary, options)
126
+ prediction, metrics, confidence = _run_rounds(
127
+ agents,
128
+ rounds,
129
+ source,
130
+ target,
131
+ context or "",
132
+ reasoning,
133
+ )
134
+ accept = True
135
+ if active_learning:
136
+ accept = active_learning_score(confidence, f1_score) > 0
137
+ return prediction, metrics, confidence, accept
138
+
139
+
140
+ class SingleModelInference(InferenceStrategy):
141
+ """One model, one round."""
142
+
143
+ def agents(self, backend, model_name, agent_configs):
144
+ """Return one agent for ``backend``/``model_name`` and one round."""
145
+ return create_agents(1, backend, model_name), 1
146
+
147
+
148
+ class DebateInference(InferenceStrategy):
149
+ """One agent per config, across the configured number of rounds.
150
+
151
+ Args:
152
+ n_rounds (int): Debate rounds.
153
+ """
154
+
155
+ def __init__(self, n_rounds: int):
156
+ self.n_rounds = n_rounds
157
+
158
+ def agents(self, backend, model_name, agent_configs):
159
+ """Return one agent per entry in ``agent_configs`` and ``n_rounds``."""
160
+ return create_agents_from_config(agent_configs), self.n_rounds
161
+
162
+
163
+ def infer(
164
+ debate: bool,
165
+ backend: str,
166
+ model_name: str,
167
+ source_term: Tuple[str, str],
168
+ target_term: Tuple[str, str],
169
+ pcmaps: Dict[str, Dict[str, set]],
170
+ dictionary: Dict[str, Any],
171
+ options: List[str],
172
+ embed_model=None,
173
+ store=None,
174
+ agent_configs: List[Dict[str, str]] = None,
175
+ n_rounds: int = 3,
176
+ dropout: float = 0.5,
177
+ reasoning: bool = False,
178
+ active_learning: bool = False,
179
+ f1_score: float = 0.0,
180
+ bisim: bool = False,
181
+ context: str = None,
182
+ ):
183
+ """Run single-model or debate inference for one alignment pair.
184
+
185
+ ``embed_model``, ``store``, ``dropout`` and ``bisim`` are accepted so
186
+ existing callers keep working. Context is built by the matching run and
187
+ passed in. See :meth:`InferenceStrategy.run` for the other arguments.
188
+
189
+ Args:
190
+ debate (bool): Use :class:`DebateInference` instead of
191
+ :class:`SingleModelInference`.
192
+
193
+ Returns:
194
+ tuple: ``(prediction, metrics, confidence, accept)``.
195
+ """
196
+ strategy = DebateInference(n_rounds) if debate else SingleModelInference()
197
+ return strategy.run(
198
+ backend=backend,
199
+ model_name=model_name,
200
+ source_term=source_term,
201
+ target_term=target_term,
202
+ pcmaps=pcmaps,
203
+ dictionary=dictionary,
204
+ options=options,
205
+ agent_configs=agent_configs or [],
206
+ n_rounds=n_rounds,
207
+ reasoning=reasoning,
208
+ active_learning=active_learning,
209
+ f1_score=f1_score,
210
+ context=context,
211
+ )