extend query example test cases ..

cover the range between simple and complex queries
biocypher · Feb 9, 2024 · 1b871b5 · 1b871b5
1 parent 1744409
commit 1b871b5
Show file tree

Hide file tree

Showing 15 changed files with 188 additions and 793 deletions.
diff --git a/benchmark/conftest.py b/benchmark/conftest.py
@@ -20,119 +20,119 @@
 # which models should be benchmarked?
 OPENAI_MODEL_NAMES = [
     "gpt-3.5-turbo",
-    "gpt-4",
+    # "gpt-4",
 ]
 
 XINFERENCE_MODELS = {
-    "llama-2-chat": {
-        "model_size_in_billions": [
-            7,
-            13,
-            70,
-        ],
-        "model_format": "ggufv2",
-        "quantization": [
-            "Q2_K",
-            # "Q3_K_S",
-            "Q3_K_M",
-            # "Q3_K_L",
-            # "Q4_0",
-            # "Q4_K_S",
-            "Q4_K_M",
-            # "Q5_0",
-            # "Q5_K_S",
-            "Q5_K_M",
-            # "Q6_K",
-            "Q8_0",
-        ],
-    },
-    "code-llama-instruct": {
-        "model_size_in_billions": [
-            7,
-            13,
-            34,
-        ],
-        "model_format": "ggufv2",
-        "quantization": [
-            "Q2_K",
-            # "Q3_K_L",
-            "Q3_K_M",
-            # "Q3_K_S",
-            # "Q4_0",
-            "Q4_K_M",
-            # "Q4_K_S",
-            # "Q5_0",
-            "Q5_K_M",
-            # "Q5_K_S",
-            # "Q6_K",
-            "Q8_0",
-        ],
-    },
-    "mixtral-instruct-v0.1": {
-        "model_size_in_billions": [
-            "46_7",
-        ],
-        "model_format": "ggufv2",
-        "quantization": [
-            "Q2_K",
-            "Q3_K_M",
-            # "Q4_0",
-            "Q4_K_M",
-            # "Q5_0",
-            "Q5_K_M",
-            # "Q6_K",
-            "Q8_0",
-        ],
-    },
-    "openhermes-2.5": {
-        "model_size_in_billions": [
-            7,
-        ],
-        "model_format": "ggufv2",
-        "quantization": [
-            "Q2_K",
-            # "Q3_K_S",
-            "Q3_K_M",
-            # "Q3_K_L",
-            # "Q4_0",
-            # "Q4_K_S",
-            "Q4_K_M",
-            # "Q5_0",
-            # "Q5_K_S",
-            "Q5_K_M",
-            # "Q6_K",
-            "Q8_0",
-        ],
-    },
-    "chatglm3": {
-        "model_size_in_billions": [
-            6,
-        ],
-        "model_format": "ggmlv3",
-        "quantization": [
-            "q4_0",
-        ],
-    },
-    "mistral-instruct-v0.2": {
-        "model_size_in_billions": [
-            7,
-        ],
-        "model_format": "ggufv2",
-        "quantization": [
-            "Q2_K",
-            # "Q3_K_S",
-            "Q3_K_M",
-            # "Q3_K_L",
-            # "Q4_0",
-            # "Q4_K_S",
-            "Q4_K_M",
-            # "Q5_0",
-            # "Q5_K_S",
-            "Q5_K_M",
-            # "Q6_K",
-            "Q8_0",
-        ],
-    },
+    # "llama-2-chat": {
+    #     "model_size_in_billions": [
+    #         7,
+    #         13,
+    #         70,
+    #     ],
+    #     "model_format": "ggufv2",
+    #     "quantization": [
+    #         "Q2_K",
+    #         # "Q3_K_S",
+    #         "Q3_K_M",
+    #         # "Q3_K_L",
+    #         # "Q4_0",
+    #         # "Q4_K_S",
+    #         "Q4_K_M",
+    #         # "Q5_0",
+    #         # "Q5_K_S",
+    #         "Q5_K_M",
+    #         # "Q6_K",
+    #         "Q8_0",
+    #     ],
+    # },
+    # "code-llama-instruct": {
+    #     "model_size_in_billions": [
+    #         7,
+    #         13,
+    #         34,
+    #     ],
+    #     "model_format": "ggufv2",
+    #     "quantization": [
+    #         "Q2_K",
+    #         # "Q3_K_L",
+    #         "Q3_K_M",
+    #         # "Q3_K_S",
+    #         # "Q4_0",
+    #         "Q4_K_M",
+    #         # "Q4_K_S",
+    #         # "Q5_0",
+    #         "Q5_K_M",
+    #         # "Q5_K_S",
+    #         # "Q6_K",
+    #         "Q8_0",
+    #     ],
+    # },
+    # "mixtral-instruct-v0.1": {
+    #     "model_size_in_billions": [
+    #         "46_7",
+    #     ],
+    #     "model_format": "ggufv2",
+    #     "quantization": [
+    #         "Q2_K",
+    #         "Q3_K_M",
+    #         # "Q4_0",
+    #         "Q4_K_M",
+    #         # "Q5_0",
+    #         "Q5_K_M",
+    #         # "Q6_K",
+    #         "Q8_0",
+    #     ],
+    # },
+    # "openhermes-2.5": {
+    #     "model_size_in_billions": [
+    #         7,
+    #     ],
+    #     "model_format": "ggufv2",
+    #     "quantization": [
+    #         "Q2_K",
+    #         # "Q3_K_S",
+    #         "Q3_K_M",
+    #         # "Q3_K_L",
+    #         # "Q4_0",
+    #         # "Q4_K_S",
+    #         "Q4_K_M",
+    #         # "Q5_0",
+    #         # "Q5_K_S",
+    #         "Q5_K_M",
+    #         # "Q6_K",
+    #         "Q8_0",
+    #     ],
+    # },
+    # "chatglm3": {
+    #     "model_size_in_billions": [
+    #         6,
+    #     ],
+    #     "model_format": "ggmlv3",
+    #     "quantization": [
+    #         "q4_0",
+    #     ],
+    # },
+    # "mistral-instruct-v0.2": {
+    #     "model_size_in_billions": [
+    #         7,
+    #     ],
+    #     "model_format": "ggufv2",
+    #     "quantization": [
+    #         "Q2_K",
+    #         # "Q3_K_S",
+    #         "Q3_K_M",
+    #         # "Q3_K_L",
+    #         # "Q4_0",
+    #         # "Q4_K_S",
+    #         "Q4_K_M",
+    #         # "Q5_0",
+    #         # "Q5_K_S",
+    #         "Q5_K_M",
+    #         # "Q6_K",
+    #         "Q8_0",
+    #     ],
+    # },
 }
 
 # create concrete benchmark list by concatenating all combinations of model

diff --git a/benchmark/data/benchmark_data.yaml b/benchmark/data/benchmark_data.yaml
@@ -9,6 +9,24 @@
 # If any input is a dictionary itself, it will be expanded into separate test
 # cases, using the top-level key to create a concatenated test case purpose.
 biocypher_query_generation:
+  - case: simple
+    input:
+      kg_path: test_schema_info.yaml
+      prompt: What is the name of the disease with ICD10 code 'E10'?
+    expected:
+      entities: ["Disease"]
+      relationships: []
+      relationship_labels: {}
+      properties:
+        Disease: ["name", "ICD10"]
+      parts_of_query:
+        [
+          "^MATCH",
+          "RETURN",
+          "([a-zA-Z]*:Disease)",
+          "WHERE [a-zA-Z]*\\.ICD10|{ICD10:",
+        ]
+
   - case: single_word
     input:
       kg_path: test_schema_info.yaml
@@ -21,25 +39,18 @@ biocypher_query_generation:
           source: Disease
           target: ["Protein", "Gene"]
       properties:
-        Disease:
-          name: {}
-          ICD10: {}
-          DSM5: {}
-        Gene:
-          name: true
-        GeneToPhenotypeAssociation:
-          score: true
-          source: true
-          evidence: true
+        Disease: ["name", "ICD10", "DSM5"]
+        Gene: ["id", "name"]
+        GeneToPhenotypeAssociation: ["score", "source", "evidence"]
       parts_of_query:
         [
           "^MATCH",
           "RETURN",
           "([a-zA-Z]*:Gene)",
           "([a-zA-Z]*:Disease)",
-          "[a-zA-Z]*:PERTURBED]",
+          "([a-zA-Z]*)?:PERTURBED]",
           "[Mm]ucoviscidosis",
-          "\\(([a-zA-Z]*:Gene)\\)<-[[a-zA-Z]*:PERTURBED]-|-[[a-zA-Z]*:PERTURBED]->\\(([a-zA-Z]*:Gene)\\)",
+          "\\(.*:Gene\\)<-\\[:PERTURBED\\]-|-\\[:PERTURBED\\]->\\(.*:Gene\\)",
           "WHERE [a-zA-Z]*\\.name|{name:",
         ]
 
@@ -53,24 +64,57 @@ biocypher_query_generation:
       relationship_labels:
         GENE_EXPRESSED_IN_CELL_TYPE:
           source: Gene
-          target: ["CellType"]
+          target: CellType
       properties:
-        CellType:
-          cell_type_name: {}
+        CellType: cell_type_name
         Gene: ["id", "name"]
-        GeneExpressedInCellType: ["expression_level"]
+        GeneExpressedInCellType: expression_level
       parts_of_query:
         [
           "^MATCH",
           "RETURN",
           "([a-zA-Z]*:Gene)",
           "([a-zA-Z]*:CellType.*)",
-          "[a-zA-Z]*:GENE_EXPRESSED_IN_CELL_TYPE]",
+          "([a-zA-Z]*)?:GENE_EXPRESSED_IN_CELL_TYPE]",
           "[Ff]ibroblast",
-          "\\(([a-zA-Z]*:Gene)\\)-[[a-zA-Z]*:GENE_EXPRESSED_IN_CELL_TYPE]->\\(([a-zA-Z]*:CellType.*)\\)|\\(([a-zA-Z]*:CellType.*)\\)<-[[a-zA-Z]*:GENE_EXPRESSED_IN_CELL_TYPE]-\\(([a-zA-Z]*:Gene)\\)",
+          "\\(.*:Gene\\)-\\[:GENE_EXPRESSED_IN_CELL_TYPE\\]->\\(.*:CellType.*\\)|\\(.*:CellType.*\\)<-\\[:GENE_EXPRESSED_IN_CELL_TYPE\\]-\\(.*:Gene\\)",
           "WHERE [a-zA-Z]*\\.cell_type_name|{cell_type_name:",
         ]
 
+  - case: complex
+    input:
+      kg_path: test_schema_info.yaml
+      prompt: Which proteins are associated with the disease having ICD10 code 'E10', what are their scores, and what is the gene related to these proteins?
+    expected:
+      entities: ["Protein", "Disease", "Gene"]
+      relationships: ["GeneToPhenotypeAssociation", "GeneToProteinAssociation"]
+      relationship_labels:
+        PERTURBED_IN:
+          source: Protein
+          target: Disease
+        GeneToProteinAssociation:
+          source: Gene
+          target: Protein
+      properties:
+        Disease: ["name", "ICD10"]
+        Protein: ["name", "score"]
+        Gene: ["name", "score"]
+        GeneToPhenotypeAssociation: ["score", "source", "evidence"]
+        GeneToProteinAssociation: ["score"]
+      parts_of_query:
+        [
+          "^MATCH",
+          "RETURN",
+          "([a-zA-Z]*:Protein)",
+          "([a-zA-Z]*:Disease)",
+          "([a-zA-Z]*:Gene)",
+          "([a-zA-Z]*)?:PERTURBED_IN]",
+          "([a-zA-Z]*)?:GeneToProteinAssociation]",
+          "WHERE [a-zA-Z]*\\.ICD10|{ICD10:",
+          "\\)<-\\[([a-zA-Z]*)?:PERTURBED_IN]-|-\\[([a-zA-Z]*)?:PERTURBED_IN]->\\(([a-zA-Z]*:Disease)",
+          "\\(([a-zA-Z]*(:Protein)?)\\)<-\\[([a-zA-Z]*)?:GeneToProteinAssociation]-|-\\[([a-zA-Z]*)?:GeneToProteinAssociation]->\\(([a-zA-Z]*(:Protein)?)\\)",
+        ]
+
 rag_interpretation:
   - case: explicit_relevance_no
     input:

diff --git a/benchmark/data/test_schema_info.yaml b/benchmark/data/test_schema_info.yaml
@@ -65,6 +65,8 @@ post translational interaction:
   represented_as: node
   label_as_edge: INTERACTS_POST_TRANSLATIONAL
   input_label: post_translational
+  source: protein
+  target: protein
 
 phosphorylation:
   is_a: post translational interaction
@@ -102,6 +104,18 @@ gene to disease association:
     source: str
     evidence: str
 
+gene to protein association:
+  represented_as: edge
+  is_relationship: true
+  present_in_knowledge_graph: true
+  input_label: gene_protein
+  source: gene
+  target: protein
+  properties:
+    score: float
+    source: str
+    evidence: str
+
 gene expressed in cell type:
   represented_as: edge
   is_a: gene to expression site association # biolink