Adding generate-extract command, 158. Add cell type templates #159

This PR does two things:

- Add a combined generate-extract command, fixes #158
- Adds cell type templates, fixes #159

src/ontogpt/cli.py

@@ -157,6 +157,11 @@ def write_extraction(
     default="yaml",
     help="Output format.",
 )
+auto_prefix_option = click.option(
+    "--auto-prefix",
+    default="AUTO",
+    help="Prefix to use for auto-generated classes. Default is AUTO.",
+)
 
 
 @click.group()
@@ -201,11 +206,7 @@ def main(verbose: int, quiet: bool, cache_db: str, skip_annotator):
 @click.option("--dictionary")
 @output_format_options
 @use_textract_options
-@click.option(
-    "--auto-prefix",
-    default="AUTO",
-    help="Prefix to use for auto-generated classes. Default is AUTO.",
-)
+@auto_prefix_option
 @click.option(
     "--set-slot-value",
     "-S",
@@ -304,6 +305,23 @@ def extract(
     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.argument("entity")
+def generate_extract(entity, template, output, output_format, **kwargs):
+    """Generate text using GPT and then extract knowledge from it."""
+    logging.info(f"Creating for {template}")
+    ke = SPIRESEngine(template, **kwargs)
+    logging.debug(f"Input entity: {entity}")
+    results = ke.generate_and_extract(entity)
+    write_extraction(results, output, output_format)
+
+
 @main.command()
 @template_option
 @model_option

src/ontogpt/engines/knowledge_engine.py

@@ -155,7 +155,11 @@ def __post_init__(self):
             self.mappers = [get_adapter("translator:")]
 
         self.set_up_client()
-        self.encoding = tiktoken.encoding_for_model(self.client.model)
+        try:
+            self.encoding = tiktoken.encoding_for_model(self.client.model)
+        except KeyError:
+            self.encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
+            logger.error(f"Could not find encoding for model {self.client.model}")
 
     def set_api_key(self, key: str):
         self.api_key = key

src/ontogpt/engines/spires_engine.py

@@ -92,10 +92,18 @@ def extract_from_text(
             named_entities=self.named_entities,
         )
 
+
     def _extract_from_text_to_dict(self, text: str, cls: ClassDefinition = None) -> RESPONSE_DICT:
         raw_text = self._raw_extract(text, cls)
         return self._parse_response_to_dict(raw_text, cls)
 
+    def generate_and_extract(
+            self, entity: str, **kwargs
+    ) -> ExtractionResult:
+        prompt = f"Generate a comprehensive description of {entity}.\n"
+        payload = self.client.complete(prompt)
+        return self.extract_from_text(payload, **kwargs)
+
     def generalize(
         self, object: Union[pydantic.BaseModel, dict], examples: List[EXAMPLE]
     ) -> ExtractionResult:

src/ontogpt/templates/cell_type.py

@@ -0,0 +1,203 @@
+from __future__ import annotations
+from datetime import datetime, date
+from enum import Enum
+from typing import List, Dict, Optional, Any, Union, Literal
+from pydantic import BaseModel as BaseModel, Field
+from linkml_runtime.linkml_model import Decimal
+
+metamodel_version = "None"
+version = "None"
+
+class WeakRefShimBaseModel(BaseModel):
+   __slots__ = '__weakref__'
+
+class ConfiguredBaseModel(WeakRefShimBaseModel,
+                validate_assignment = True, 
+                validate_all = True, 
+                underscore_attrs_are_private = True, 
+                extra = 'forbid', 
+                arbitrary_types_allowed = True):
+    pass                    
+
+
+class BrainRegionIdentifier(str, Enum):
+
+
+    dummy = "dummy"
+
+
+class NullDataOptions(str, Enum):
+
+    UNSPECIFIED_METHOD_OF_ADMINISTRATION = "UNSPECIFIED_METHOD_OF_ADMINISTRATION"
+    NOT_APPLICABLE = "NOT_APPLICABLE"
+    NOT_MENTIONED = "NOT_MENTIONED"
+
+
+
+class CellTypeDocument(ConfiguredBaseModel):
+
+    cell_type: Optional[str] = Field(None, description="""the name of the cell type described""")
+    range: Optional[str] = Field(None)
+    parents: Optional[List[str]] = Field(default_factory=list, description="""categorization""")
+    subtypes: Optional[List[str]] = Field(default_factory=list)
+    localizations: Optional[List[str]] = Field(default_factory=list)
+    genes: Optional[List[str]] = Field(default_factory=list)
+    diseases: Optional[List[str]] = Field(default_factory=list)
+
+
+
+class InterneuronDocument(CellTypeDocument):
+
+    projects_to_or_from: Optional[List[str]] = Field(default_factory=list, description="""Brain structures from which this cell type projects into or receives projections from""")
+    cell_type: Optional[str] = Field(None, description="""the name of the cell type described""")
+    range: Optional[str] = Field(None)
+    parents: Optional[List[str]] = Field(default_factory=list, description="""categorization""")
+    subtypes: Optional[List[str]] = Field(default_factory=list)
+    localizations: Optional[List[str]] = Field(default_factory=list)
+    genes: Optional[List[str]] = Field(default_factory=list)
+    diseases: Optional[List[str]] = Field(default_factory=list)
+
+
+
+class ExtractionResult(ConfiguredBaseModel):
+    """
+    A result of extracting knowledge on text
+    """
+    input_id: Optional[str] = Field(None)
+    input_title: Optional[str] = Field(None)
+    input_text: Optional[str] = Field(None)
+    raw_completion_output: Optional[str] = Field(None)
+    prompt: Optional[str] = Field(None)
+    extracted_object: Optional[Any] = Field(None, description="""The complex objects extracted from the text""")
+    named_entities: Optional[List[Any]] = Field(default_factory=list, description="""Named entities extracted from the text""")
+
+
+
+class NamedEntity(ConfiguredBaseModel):
+
+    id: str = Field(None, description="""A unique identifier for the named entity""")
+    label: Optional[str] = Field(None, description="""The label (name) of the named thing""")
+
+
+
+class Gene(NamedEntity):
+
+    id: str = Field(None, description="""A unique identifier for the named entity""")
+    label: Optional[str] = Field(None, description="""The label (name) of the named thing""")
+
+
+
+class Pathway(NamedEntity):
+
+    id: str = Field(None, description="""A unique identifier for the named entity""")
+    label: Optional[str] = Field(None, description="""The label (name) of the named thing""")
+
+
+
+class AnatomicalStructure(NamedEntity):
+
+    id: str = Field(None, description="""A unique identifier for the named entity""")
+    label: Optional[str] = Field(None, description="""The label (name) of the named thing""")
+
+
+
+class BrainRegion(AnatomicalStructure):
+
+    id: str = Field(None, description="""A unique identifier for the named entity""")
+    label: Optional[str] = Field(None, description="""The label (name) of the named thing""")
+
+
+
+class CellType(NamedEntity):
+
+    id: str = Field(None, description="""A unique identifier for the named entity""")
+    label: Optional[str] = Field(None, description="""The label (name) of the named thing""")
+
+
+
+class Disease(NamedEntity):
+
+    id: str = Field(None, description="""A unique identifier for the named entity""")
+    label: Optional[str] = Field(None, description="""The label (name) of the named thing""")
+
+
+
+class Drug(NamedEntity):
+
+    id: str = Field(None, description="""A unique identifier for the named entity""")
+    label: Optional[str] = Field(None, description="""The label (name) of the named thing""")
+
+
+
+class CompoundExpression(ConfiguredBaseModel):
+
+    None
+
+
+
+class Triple(CompoundExpression):
+    """
+    Abstract parent for Relation Extraction tasks
+    """
+    subject: Optional[str] = Field(None)
+    predicate: Optional[str] = Field(None)
+    object: Optional[str] = Field(None)
+    qualifier: Optional[str] = Field(None, description="""A qualifier for the statements, e.g. \"NOT\" for negation""")
+    subject_qualifier: Optional[str] = Field(None, description="""An optional qualifier or modifier for the subject of the statement, e.g. \"high dose\" or \"intravenously administered\"""")
+    object_qualifier: Optional[str] = Field(None, description="""An optional qualifier or modifier for the object of the statement, e.g. \"severe\" or \"with additional complications\"""")
+
+
+
+class TextWithTriples(ConfiguredBaseModel):
+
+    publication: Optional[Publication] = Field(None)
+    triples: Optional[List[Triple]] = Field(default_factory=list)
+
+
+
+class RelationshipType(NamedEntity):
+
+    id: str = Field(None, description="""A unique identifier for the named entity""")
+    label: Optional[str] = Field(None, description="""The label (name) of the named thing""")
+
+
+
+class Publication(ConfiguredBaseModel):
+
+    id: Optional[str] = Field(None, description="""The publication identifier""")
+    title: Optional[str] = Field(None, description="""The title of the publication""")
+    abstract: Optional[str] = Field(None, description="""The abstract of the publication""")
+    combined_text: Optional[str] = Field(None)
+    full_text: Optional[str] = Field(None, description="""The full text of the publication""")
+
+
+
+class AnnotatorResult(ConfiguredBaseModel):
+
+    subject_text: Optional[str] = Field(None)
+    object_id: Optional[str] = Field(None)
+    object_text: Optional[str] = Field(None)
+
+
+
+
+# Update forward refs
+# see https://pydantic-docs.helpmanual.io/usage/postponed_annotations/
+CellTypeDocument.update_forward_refs()
+InterneuronDocument.update_forward_refs()
+ExtractionResult.update_forward_refs()
+NamedEntity.update_forward_refs()
+Gene.update_forward_refs()
+Pathway.update_forward_refs()
+AnatomicalStructure.update_forward_refs()
+BrainRegion.update_forward_refs()
+CellType.update_forward_refs()
+Disease.update_forward_refs()
+Drug.update_forward_refs()
+CompoundExpression.update_forward_refs()
+Triple.update_forward_refs()
+TextWithTriples.update_forward_refs()
+RelationshipType.update_forward_refs()
+Publication.update_forward_refs()
+AnnotatorResult.update_forward_refs()
+