Merge pull request #147 from monarch-initiative/prep_for_0.2.4

Prep for 0.2.4 release

Author: caufieldjh

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "curategpt"
-version = "0.2.3"
+version = "0.2.4"
 description = "CurateGPT"
 authors = ["Chris Mungall <cjmungall@lbl.gov>", "Carlo Kroll <ckroll95@gmail.com>", "Harshad Hegde <hhegde@lbl.gov>", "J. Harry Caufield <jhc@lbl.gov>"]
 license = "BSD-3"

pyproject.toml

@@ -32,7 +32,8 @@ def select_from_options_prompt(
     """
     if prompt_template is None:
         if query is None:
-            raise ValueError("Either query or prompt_template must be specified.")
+            raise ValueError(
+                "Either query or prompt_template must be specified.")
         prompt_template = "I will first give background facts, then ask a question."
         prompt_template += "Use the background fact to answer\n"
         prompt_template += "---\nBackground facts:\n"
@@ -53,7 +54,8 @@ def select_from_options_prompt(
         current_length = 0
         for obj, _, _obj_meta in kb_results:
             i += 1
-            obj_text = yaml.dump({k: v for k, v in obj.items() if v}, sort_keys=False)
+            obj_text = yaml.dump(
+                {k: v for k, v in obj.items() if v}, sort_keys=False)
             references[str(i)] = obj_text
             objects[str(i)] = obj
             if id_field and id_field in obj:
@@ -65,7 +67,8 @@ def select_from_options_prompt(
         prompt = prompt_template.format(body="".join(texts), query=query)
         logger.info(f"Prompt: {prompt}")
         estimated_length = estimate_num_tokens([prompt])
-        logger.debug(f"Max tokens {model.model_id}: {max_tokens_by_model(model.model_id)}")
+        logger.debug(
+            f"Max tokens {model.model_id}: {max_tokens_by_model(model.model_id)}")
         # TODO: use a more precise estimate of the length
         if estimated_length + 300 < max_tokens_by_model(model.model_id):
             break

src/curategpt/agents/agent_utils.py

@@ -28,4 +28,5 @@ class BaseAgent(ABC):  # noqa: B024
     """Engine performing LLM operations, including extracting from prompt responses"""
 
     def search(self):
-        raise NotImplementedError("Search method must be implemented by subclass")
+        raise NotImplementedError(
+            "Search method must be implemented by subclass")

src/curategpt/agents/base_agent.py

@@ -1,6 +1,7 @@
 """Annotation (Concept Recognition) in texts."""
 
 import logging
+import re
 from dataclasses import dataclass
 from enum import Enum
 from typing import Dict, List, Optional, Tuple
@@ -129,9 +130,6 @@ class AnnotatedText(BaseModel):
 """
 
 
-import re
-
-
 def parse_annotations(text, marker_char: str = None) -> List[CONCEPT]:
     """
     Parse annotations from text.
@@ -283,7 +281,8 @@ def ground_concept(
                 )
                 spans.append(span)
         # spans = parse_spans(response.text(), concept_dict)
-        ann = GroundingResult(input_text=text, annotated_text=response.text(), spans=spans)
+        ann = GroundingResult(
+            input_text=text, annotated_text=response.text(), spans=spans)
         return ann
 
     def annotate(
@@ -340,7 +339,8 @@ def annotate_two_pass(
                 **kwargs,
             )
             if not concept.spans:
-                logger.debug(f"Unable to ground concept {term} in category {category}")
+                logger.debug(
+                    f"Unable to ground concept {term} in category {category}")
                 continue
             main_span = concept.spans[0]
             spans.append(
@@ -376,7 +376,8 @@ def annotate_inline(
 
         logger.info(f"Anns: {anns}")
         spans = [
-            Span(text=ann[0], concept_id=ann[1], concept_label=concept_dict.get(ann[1], None))
+            Span(text=ann[0], concept_id=ann[1],
+                 concept_label=concept_dict.get(ann[1], None))
             for ann in anns
         ]
         return AnnotatedText(input_text=text, spans=spans, annotated_text=response.text())
@@ -436,7 +437,8 @@ def _label_id_pairs_prompt_section(
             if not id:
                 raise ValueError(f"Object {obj} has no ID field {id_field}")
             if not label:
-                raise ValueError(f"Object {obj} has no label field {label_field}")
+                raise ValueError(
+                    f"Object {obj} has no label field {label_field}")
             prompt += f"{label} // {id}   \n"
             concept_pairs.append((id, label))
         return concept_pairs, prompt

src/curategpt/agents/concept_recognition_agent.py

@@ -65,7 +65,8 @@ class DatabaseAugmentedStructuredExtraction(BaseAgent):
     background_document_limit: int = 3
     """Number of background documents to use. TODO: more sophisticated way to estimate."""
 
-    default_masked_fields: List[str] = field(default_factory=lambda: ["original_id"])
+    default_masked_fields: List[str] = field(
+        default_factory=lambda: ["original_id"])
 
     def extract(
         self,
@@ -135,7 +136,7 @@ def extract(
             ):
                 obj_text = obj_meta["document"]
                 # TODO: use tiktoken to estimate
-                obj_text = obj_text[0 : self.max_background_document_size]
+                obj_text = obj_text[0: self.max_background_document_size]
                 docs.append(obj_text)
         if generate_background:
             # prompt = f"Generate a comprehensive description about the {target_class} with {context_property} = {seed}"

src/curategpt/agents/dase_agent.py

@@ -67,7 +67,8 @@ class DragonAgent(BaseAgent):
     background_document_limit: int = 3
     """Number of background documents to use. TODO: more sophisticated way to estimate."""
 
-    default_masked_fields: List[str] = field(default_factory=lambda: ["original_id"])
+    default_masked_fields: List[str] = field(
+        default_factory=lambda: ["original_id"])
 
     def complete(
         self,
@@ -139,10 +140,12 @@ def generate_input_str(obj: Union[str, Dict], prefix="Structured representation
             elif isinstance(obj, str):
                 return f"{prefix} {target_class} with {context_property} = {obj}"
             else:
-                raise ValueError(f"Invalid type for obj: {type(obj)} //  {obj}")
+                raise ValueError(
+                    f"Invalid type for obj: {type(obj)} //  {obj}")
 
         annotated_examples = []
-        seed_search_term = seed if isinstance(seed, str) else yaml.safe_dump(seed, sort_keys=True)
+        seed_search_term = seed if isinstance(
+            seed, str) else yaml.safe_dump(seed, sort_keys=True)
         logger.debug(f"Searching for seed: {seed_search_term}")
         for obj, _, _obj_meta in self.knowledge_source.search(
             seed_search_term,
@@ -167,7 +170,8 @@ def generate_input_str(obj: Union[str, Dict], prefix="Structured representation
                     f"Num examples={len(annotated_examples)}"
                 )
                 continue
-            ae = AnnotatedObject(object=obj_predicted_part, annotations={"text": input_text})
+            ae = AnnotatedObject(object=obj_predicted_part,
+                                 annotations={"text": input_text})
             annotated_examples.append(ae)
         if not annotated_examples:
             logger.error(f"No suitable examples found for seed: {seed}")
@@ -181,7 +185,7 @@ def generate_input_str(obj: Union[str, Dict], prefix="Structured representation
             ):
                 obj_text = obj_meta["document"]
                 # TODO: use tiktoken to estimate
-                obj_text = obj_text[0 : self.max_background_document_size]
+                obj_text = obj_text[0: self.max_background_document_size]
                 docs.append(obj_text)
         gen_text = generate_input_str(seed)
         if generate_background:
@@ -249,7 +253,8 @@ def generate_all(
                 continue
             curr_val = obj.get(field_to_predict, None)
             if missing_only and curr_val:
-                logger.debug(f"Skipping; {field_to_predict} already present: {curr_val}")
+                logger.debug(
+                    f"Skipping; {field_to_predict} already present: {curr_val}")
                 continue
             ao = self.complete(obj, collection=collection, **kwargs)
             yield PredictedFieldValue(
@@ -286,7 +291,7 @@ def generate_queries(self, context_property="name", n=5, **kwargs) -> List[str]:
         response = self.extractor.model.prompt(prompt, system=system)
         txt = response.text()
         if "." in txt:
-            txt = txt[0 : txt.index(".")]
+            txt = txt[0: txt.index(".")]
         suggestions = txt.split(",")
         suggestions = [x.strip() for x in suggestions]
         ids_norm = [x.lower() for x in ids]
@@ -309,7 +314,8 @@ def review(
         :param obj:
         """
         if fields_to_predict and not context_property:
-            raise ValueError("context_property is required if fields_to_predict")
+            raise ValueError(
+                "context_property is required if fields_to_predict")
 
         pk_val = obj.get(primary_key, None)
 
@@ -360,7 +366,8 @@ def _obj_as_str(obj: dict) -> str:
         if not isinstance(ao.object, dict):
             logger.warning(f"Expected dict, got {ao.object}")
             if isinstance(ao.object, list):
-                logger.warning(f"Taking first element of list of len {len(ao.object)}")
+                logger.warning(
+                    f"Taking first element of list of len {len(ao.object)}")
                 ao.object = ao.object[0]
         if isinstance(ao.object, dict):
             if context_property:

src/curategpt/agents/dragon_agent.py

@@ -38,7 +38,8 @@ class EvidenceAgent(BaseAgent):
 
     chat_agent: Union[ChatAgent, BaseWrapper] = None
 
-    evidence_update_policy: EvidenceUpdatePolicyEnum = field(default=EvidenceUpdatePolicyEnum.skip)
+    evidence_update_policy: EvidenceUpdatePolicyEnum = field(
+        default=EvidenceUpdatePolicyEnum.skip)
 
     def find_evidence(self, obj: Union[str, Dict]) -> ChatResponse:
         obj_as_str = obj if isinstance(obj, str) else object_as_yaml(obj)
@@ -78,7 +79,8 @@ def find_evidence_simple(self, query: str, limit: int = 10, **kwargs) -> Optiona
             current_length = 0
             for obj, _, _obj_meta in kb_results:
                 i += 1
-                obj_text = yaml.dump({k: v for k, v in obj.items() if v}, sort_keys=False)
+                obj_text = yaml.dump(
+                    {k: v for k, v in obj.items() if v}, sort_keys=False)
                 references[str(i)] = obj_text
                 texts.append(f"## Reference\n{obj_text}")
                 current_length += len(obj_text)
@@ -97,7 +99,8 @@ def find_evidence_simple(self, query: str, limit: int = 10, **kwargs) -> Optiona
                 break
             else:
                 # remove least relevant
-                logger.debug(f"Removing least relevant of {len(kb_results)}: {kb_results[-1]}")
+                logger.debug(
+                    f"Removing least relevant of {len(kb_results)}: {kb_results[-1]}")
                 if not kb_results:
                     raise ValueError(f"Prompt too long: {prompt}.")
                 kb_results.pop()

src/curategpt/agents/evidence_agent.py

@@ -13,6 +13,7 @@
 HF_DOWNLOAD_PATH = Path(__file__).resolve().parents[4]
 HF_DOWNLOAD_PATH = HF_DOWNLOAD_PATH / "hf_download"
 
+
 @dataclass
 class HuggingFaceAgent:
 
@@ -38,13 +39,17 @@ def upload(self, objects, metadata, repo_id, private=False, **kwargs):
         metadata_file = "metadata.yaml"
 
         try:
-            df = pd.DataFrame(data=[(obj[0], obj[2]['_embeddings'], obj[2]['document']) for obj in objects])
+            df = pd.DataFrame(
+                data=[(obj[0], obj[2]['_embeddings'], obj[2]['document']) for obj in objects])
         except Exception as e:
-            raise ValueError(f"Creation of Dataframe not successful: {e}") from e
+            raise ValueError(
+                f"Creation of Dataframe not successful: {e}") from e
 
         with ExitStack() as stack:
-            tmp_parquet = stack.enter_context(tempfile.NamedTemporaryFile(suffix=".parquet", delete=True))
-            tmp_yaml = stack.enter_context(tempfile.NamedTemporaryFile(suffix=".yaml", delete=True))
+            tmp_parquet = stack.enter_context(
+                tempfile.NamedTemporaryFile(suffix=".parquet", delete=True))
+            tmp_yaml = stack.enter_context(
+                tempfile.NamedTemporaryFile(suffix=".yaml", delete=True))
 
             embedding_path = tmp_parquet.name
             metadata_path = tmp_yaml.name
@@ -56,8 +61,8 @@ def upload(self, objects, metadata, repo_id, private=False, **kwargs):
             self._create_repo(repo_id, private=private)
 
             self._upload_files(repo_id, {
-                embedding_path : repo_id + "/" + embedding_file,
-                metadata_path : repo_id + "/" + metadata_file
+                embedding_path: repo_id + "/" + embedding_file,
+                metadata_path: repo_id + "/" + metadata_file
             })
 
     def upload_duckdb(self, objects, metadata, repo_id, private=False, **kwargs):
@@ -74,13 +79,17 @@ def upload_duckdb(self, objects, metadata, repo_id, private=False, **kwargs):
         embedding_file = "embeddings.parquet"
         metadata_file = "metadata.yaml"
         try:
-            df = pd.DataFrame(data=[(obj[0], obj[2]['_embeddings'], obj[2]['documents']) for obj in objects])
+            df = pd.DataFrame(
+                data=[(obj[0], obj[2]['_embeddings'], obj[2]['documents']) for obj in objects])
         except Exception as e:
-            raise ValueError(f"Creation of Dataframe not successful: {e}") from e
+            raise ValueError(
+                f"Creation of Dataframe not successful: {e}") from e
 
         with ExitStack() as stack:
-            tmp_parquet = stack.enter_context(tempfile.NamedTemporaryFile(suffix=".parquet", delete=True))
-            tmp_yaml = stack.enter_context(tempfile.NamedTemporaryFile(suffix=".yaml", delete=True))
+            tmp_parquet = stack.enter_context(
+                tempfile.NamedTemporaryFile(suffix=".parquet", delete=True))
+            tmp_yaml = stack.enter_context(
+                tempfile.NamedTemporaryFile(suffix=".yaml", delete=True))
 
             embedding_path = tmp_parquet.name
             metadata_path = tmp_yaml.name
@@ -92,8 +101,8 @@ def upload_duckdb(self, objects, metadata, repo_id, private=False, **kwargs):
             self._create_repo(repo_id, private=private)
 
             self._upload_files(repo_id, {
-                embedding_path : repo_id + "/" + embedding_file,
-                metadata_path : repo_id + "/" + metadata_file
+                embedding_path: repo_id + "/" + embedding_file,
+                metadata_path: repo_id + "/" + metadata_file
             })
 
     def _create_repo(self, repo_id: str, private: bool = False):
@@ -104,13 +113,15 @@ def _create_repo(self, repo_id: str, private: bool = False):
         :param private: Whether the repository is private.
         """
         try:
-            create_repo(repo_id=repo_id, token=self.token, repo_type="dataset", private=private)
-            logger.info(f"Repository {repo_id} created successfully on Hugging Face.")
+            create_repo(repo_id=repo_id, token=self.token,
+                        repo_type="dataset", private=private)
+            logger.info(
+                f"Repository {repo_id} created successfully on Hugging Face.")
         except Exception as e:
-            logger.error(f"Failed to create repository {repo_id} on Hugging Face: {e}")
+            logger.error(
+                f"Failed to create repository {repo_id} on Hugging Face: {e}")
             raise
 
-
     def _upload_files(self, repo_id: str, files: Dict[str, str]):
         """
         Upload files to a Hugging Face repository.
@@ -126,9 +137,11 @@ def _upload_files(self, repo_id: str, files: Dict[str, str]):
                     repo_id=repo_id,
                     repo_type="dataset",
                 )
-                logger.info(f"Uploaded {local_path} to {repo_path} in {repo_id}")
+                logger.info(
+                    f"Uploaded {local_path} to {repo_path} in {repo_id}")
         except Exception as e:
-            logger.error(f"Failed to upload files to {repo_id} on Hugging Face: {e}")
+            logger.error(
+                f"Failed to upload files to {repo_id} on Hugging Face: {e}")
             raise
 
     def cached_download(
@@ -145,7 +158,3 @@ def cached_download(
         )
 
         return download_path
-
-
-
-