Improvements to iterative generate-extract. (#166)

pyproject.toml

@@ -14,13 +14,13 @@ oaklib = ">=0.5.12"
 gilda = ">=1.0.0"
 jsonlines = ">=3.1.0"
 python-multipart = "^0.0.5"
-linkml-owl = ">=0.2.7"
+linkml-owl = "^0.2.8"
 beautifulsoup4 = ">=4.11.1"
 eutils = ">=0.6.0"
 class-resolver = ">=0.4.2"
 inflect = ">=6.0.2"
 bioc = ">=2.0.post5"
-linkml = ">=1.4.10"
+linkml = "^1.5.6"
 wikipedia = ">=1.4.0"
 tiktoken = ">=0.3.3"
 airium = ">=0.2.5"

src/ontogpt/cli.py

@@ -22,16 +22,16 @@
 from sssom.util import to_mapping_set_dataframe
 
 import ontogpt.ontex.extractor as extractor
-from ontogpt import MODELS, DEFAULT_MODEL, __version__
+from ontogpt import DEFAULT_MODEL, MODELS, __version__
 from ontogpt.clients import OpenAIClient
 from ontogpt.clients.pubmed_client import PubmedClient
 from ontogpt.clients.soup_client import SoupClient
 from ontogpt.clients.wikipedia_client import WikipediaClient
 from ontogpt.engines import create_engine
 from ontogpt.engines.embedding_similarity_engine import SimilarityEngine
 from ontogpt.engines.enrichment import EnrichmentEngine
-from ontogpt.engines.ggml_engine import GGMLEngine
 from ontogpt.engines.generic_engine import GenericEngine, QuestionCollection
+from ontogpt.engines.ggml_engine import GGMLEngine
 from ontogpt.engines.halo_engine import HALOEngine
 from ontogpt.engines.hfhub_engine import HFHubEngine
 from ontogpt.engines.knowledge_engine import KnowledgeEngine
@@ -318,7 +318,50 @@ def generate_extract(entity, template, output, output_format, **kwargs):
     ke = SPIRESEngine(template, **kwargs)
     logging.debug(f"Input entity: {entity}")
     results = ke.generate_and_extract(entity)
-    write_extraction(results, output, output_format)
+    write_extraction(results, output, output_format, ke)
+
+
+@main.command()
+@template_option
+@model_option
+@recurse_option
+@output_option_wb
+@output_format_options
+@auto_prefix_option
+@click.option("--ontology", "-r", help="Ontology to use; use oaklib selector path")
+@click.option("--max-iterations", "-M", default=10, type=click.INT)
+@click.option("--iteration-slot", "-I", multiple=True, help="Slots to iterate over")
+@click.option("--db", "-D", help="Where the resulting yaml database is stored")
+@click.option(
+    "--clear/--no-clear", default=False, show_default=True, help="Clear the db before starting"
+)
+@click.argument("entity")
+def iteratively_generate_extract(
+    entity,
+    template,
+    output,
+    output_format,
+    db,
+    iteration_slot,
+    max_iterations,
+    clear,
+    ontology,
+    **kwargs,
+):
+    """Iterate through generate-extract."""
+    logging.info(f"Creating for {template}")
+    ke = SPIRESEngine(template, **kwargs)
+    logging.debug(f"Input entity: {entity}")
+    adapter = get_adapter(ontology)
+    for results in ke.iteratively_generate_and_extract(
+        entity,
+        db,
+        iteration_slots=list(iteration_slot),
+        max_iterations=max_iterations,
+        adapter=adapter,
+        clear=clear,
+    ):
+        write_extraction(results, output, output_format)
 
 
 @main.command()