Merge pull request #28 from asapdiscovery/stricter_endpoints

Make LLM a bit stricter on endpoints

falcbot/falcbot.py

@@ -80,9 +80,12 @@ def make_pic50_pred(event, say, context, logger):
         say(f"Invalid SMILES {smiles}, unable to proceed")
         return
 
+    targets = ASAPMLModelRegistry.get_targets_with_models()
+    # filter out None values
+    targets = [t for t in targets if t is not None]
     if not target in ASAPMLModelRegistry.get_targets_with_models():
         say(
-            f"Invalid target {target}, not in: {ASAPMLModelRegistry.get_targets_with_models()}; unable to proceed"
+            f"Invalid target {target}, not in: {targets}; unable to proceed"
         )
         return
 

falcbot/llm.py

@@ -27,18 +27,18 @@ class IsMLQuery(BaseModel):
 def _and_join(lst):
     return " and ".join(lst)
 
-_base_ml_prompt_template = "You are an expert scientist, parse the following making sure all SMILES strings are represented exactly as in the input: Be very careful and use only SMILES already in the prompt. Allowed variables for target are {targets} and for property are {properties} : {query}"
+_base_ml_prompt_template = "You are an expert scientist, parse the following making sure all SMILES strings are represented exactly as in the input: Be very careful and use only SMILES already in the prompt. Allowed variables for target are {targets} and for property are {properties}. If properties do not appear close to the allowed values, use None. For targets you can be more generous: {query}"
 
 def _make_ml_prompt_template() -> PromptTemplate:
     """
     Create a prompt template for the ASAPMLModelQuery model
     """
     # join to make a string with "and" between each
     targets_w_models = ASAPMLModelRegistry.get_targets_with_models()
-    # filter out None values
-    targets_w_models = [t for t in targets_w_models if t is not None]        
+    targets_w_models = [t for t in targets_w_models if t is not None] + ["None"]
+    properties = ASAPMLModelRegistry.get_endpoints() + ["None"]       
     target_str = _and_join(targets_w_models)
-    properties = _and_join(ASAPMLModelRegistry.get_endpoints())
+    properties = _and_join(properties)
     pt = PromptTemplate(_base_ml_prompt_template)
     formatted =  pt.partial_format(targets=target_str, properties=properties)
     return formatted