Upload the second version

esgen/functions.py

@@ -9,16 +9,17 @@ from esgen.verbalizer import chatbot_verbaliser, init_comment_verbaliser
 from esgen.utils import wikidata_api_search
 
 
-def entity_schema_generation(name, class_id, threshold):
+def entity_schema_generation(name, class_id, threshold, property_types):
     """
 
     :param name:
     :param class_id:
     :param threshold:
+    :param property_types:
     :return:
     """
     try:
-        es = EntitySchema(name, class_id, threshold / 100)
+        es = EntitySchema(name, class_id, threshold / 100, property_types)
         es.properties, es.pending = es.generate_initial_schema()
         prop = es.pending[0]
         prop["examples"] = collect_examples(es.class_id, prop["id"])
@@ -91,7 +92,14 @@ def input_no(es_json, history):
         es.load_es_json(es_json)
         return str(es), es.__json__(), history
     # add the pending property to rejected properties
-    es_json["rejected"].append(es_json["pending"].pop(0))
+    rejected_prop = es_json["pending"].pop(0)
+    constraint_json = {
+        "prop": "wdt:" + rejected_prop["id"],
+        "value": ".",  # with any value
+        "cardinality": "{0}",  # has no matching statements
+        "comment": init_comment_verbaliser(rejected_prop["label"])
+    }
+    es_json["rejected"].append(constraint_json)
     es = EntitySchema(name="", class_id="")
     es.load_es_json(es_json)
     try:

esgen/model.py

@@ -50,6 +50,7 @@ class EntitySchema:
             threshold=0.5,
             extra=False,
             prefixes=None,
+            property_types=None,
     ):
         """
 
@@ -58,6 +59,7 @@ class EntitySchema:
         :param extra: may have other statements
         :param prefixes:
         :param threshold:
+        :param property_types:
         """
         self.name = name
         self.class_id = class_id
@@ -72,6 +74,8 @@ class EntitySchema:
             self.prefixes = default_prefixes
         else:
             self.prefixes = prefixes
+        if property_types is None:
+            self.property_types = ["wikibase:WikibaseItem"]
         self.properties = list()
         self.pending = list()
         self.optional = list()
@@ -82,12 +86,12 @@ class EntitySchema:
         self.parent_relation = None  # relation connect es to its parent es, usually used for sub es
 
     def __str__(self):
-        prefix_list = [prefix_verbaliser(abbr, url)for abbr, url in self.prefixes.items()]
+        prefix_list = [prefix_verbaliser(abbr, url) for abbr, url in self.prefixes.items()]
         prefix_str = "".join(prefix_list)
 
         start_str = f"start = @<{self.name}>\n"
 
-        constraint_list = self.properties + self.optional
+        constraint_list = self.properties + self.optional + self.rejected
         constraint_str = value_shape_verbaliser(self.name, constraint_list, self.extra)
 
         shex_script = '\n'.join([prefix_str, start_str, constraint_str])
@@ -96,7 +100,12 @@ class EntitySchema:
 
     def generate_initial_schema(self):
         schema_components, pending = list(), list()
-        properties = filter_properties(class_id=self.class_id, num_instances=NUM_ITEMS, threshold=self.threshold)
+        properties = filter_properties(
+            class_id=self.class_id,
+            num_instances=NUM_ITEMS,
+            threshold=self.threshold,
+            property_types=self.property_types
+        )
         # print(properties)
         for prop_id, prop in properties.items():
             if prop["frequency_level"] == "frequent":
@@ -197,11 +206,16 @@ class NestedEntitySchema(EntitySchema):
 
         start_str = f"\nstart = @<{self.name}>\n"
 
-        constraint_list = self.properties + self.optional
+        self.sort_constraints()  # sort constraints before verbalization
+        constraint_list = self.properties + self.optional + self.rejected
         constraint_str = value_shape_verbaliser(self.name, constraint_list, self.extra)
 
         value_shapes_str = "\n".join(
-            value_shape_verbaliser(value_shape_name, value_shape.properties + value_shape.optional, extra=False)
+            value_shape_verbaliser(
+                value_shape_name,
+                value_shape.properties + value_shape.optional + value_shape.rejected,
+                extra=False
+            )
             for value_shape_name, value_shape in self.value_shapes.items()
         )
 
@@ -347,4 +361,11 @@ class NestedEntitySchema(EntitySchema):
                         self.properties.append(constraint)
 
     def sort_constraints(self):
-        pass
+        # rule 1: keep "instance of" property in the first place
+        sorted_properties = list()
+        for constraint in self.properties:
+            if constraint.prop == "wdt:P31":
+                sorted_properties.insert(0, constraint)
+            else:
+                sorted_properties.append(constraint)
+        self.properties = sorted_properties

esgen/queries.py

@@ -55,10 +55,11 @@ def collect_instances(class_id: str, instances_num: int) -> list:
     return instances
 
 
-def collect_properties(instances: list) -> dict:
+def collect_properties(instances: list, property_types: list) -> dict:
     """
 
     :param instances:
+    :param property_types:
     :return:
     """
     # properties = []
@@ -69,8 +70,9 @@ def collect_properties(instances: list) -> dict:
           wd:%s ?property ?value .
           ?prop wikibase:directClaim ?property .
           ?prop wikibase:propertyType ?propType .
+          FILTER ( ?propType IN ( %s ) )
           SERVICE wikibase:label { bd:serviceParam wikibase:language "en" }
-        }""" % instance
+        }""" % (instance, ",".join(property_types))
         results = get_sparql_results(query)
         for result in results["results"]["bindings"]:
             # property_count = int(result["count"]["value"])
@@ -94,18 +96,19 @@ def collect_properties(instances: list) -> dict:
     return properties
 
 
-def filter_properties(class_id: str, num_instances: int, threshold: float) -> dict:
+def filter_properties(class_id: str, num_instances: int, threshold: float, property_types: list) -> dict:
     """
     filter properties based on the threshold
 
     :param class_id:
     :param num_instances:
     :param threshold:
+    :param property_types:
     :return:
     """
     # instance_count = count_instances(class_id)
     instances = collect_instances(class_id, num_instances)
-    properties = collect_properties(instances)
+    properties = collect_properties(instances, property_types)
     for key, value in properties.items():
         # prop_count = count_properties(instances, prop["property"])
         # prop["count"] = prop_count

esgen/verbalizer.py

@@ -1,5 +1,8 @@
-def init_comment_verbaliser(prop_label, freq):
-    comment = f"{prop_label}, {freq * 100:.2f}%"
+def init_comment_verbaliser(prop_label, freq=None):
+    if freq:
+        comment = f"{prop_label}, {freq * 100:.2f}%"
+    else:
+        comment = f"{prop_label}"
     return comment