Merge pull request #223 from adhardydm/use-custom-model-file

Use custom model file

Author: nielstron

.travis.yml

@@ -50,6 +50,7 @@ script:
   - coverage run -a --source=quantulum3 setup.py test -s quantulum3.tests.test_classifier.ClassifierBuild
   - coverage run -a --source=quantulum3 setup.py test -s quantulum3.tests.test_classifier.ClassifierTest
   - coverage run -a --source=quantulum3 setup.py test -s quantulum3.tests.test_scripts.TrainScriptTest
+  - coverage run -a --source=quantulum3 setup.py test -s quantulum3.tests.test_load.TestCached
 
 after_success:
   - coverage report

README.md

@@ -184,6 +184,14 @@ converted to a `.magnitude` file on-the-run. Check out
 and [Magnitude](https://github.com/plasticityai/magnitude) for more information.
 
 
+To use your custom model, pass the path to the trained model file to the
+parser:
+
+```pyton
+parser = Parser.parse(classifier_path="path/to/model")
+```
+
+
 Manipulation
 ------------
 

quantulum3/_lang/en_US/parser.py

@@ -234,7 +234,9 @@ def parse_unit(_, unit, slash):
 
 
 ###############################################################################
-def build_quantity(orig_text, text, item, values, unit, surface, span, uncert):
+def build_quantity(
+    orig_text, text, item, values, unit, surface, span, uncert, classifier_path=None
+):
     """
     Build a Quantity object out of extracted information.
     """
@@ -425,7 +427,7 @@ def build_quantity(orig_text, text, item, values, unit, surface, span, uncert):
     if dimension_change:
         if unit.original_dimensions:
             unit = parser.get_unit_from_dimensions(
-                unit.original_dimensions, orig_text, lang
+                unit.original_dimensions, orig_text, lang, classifier_path
             )
         else:
             unit = load.units(lang).names["dimensionless"]

quantulum3/classifier.py

@@ -7,6 +7,7 @@
 import logging
 import multiprocessing
 import os
+import warnings
 
 import pkg_resources
 
@@ -24,6 +25,11 @@
     SGDClassifier, TfidfVectorizer = None, None
     USE_CLF = False
 
+    warnings.warn(
+        "Classifier dependencies not installed. Run pip install quantulum3[classifier] "
+        "to install them. The classifer helps to dissambiguate units."
+    )
+
 try:
     import wikipedia
 except ImportError:
@@ -203,9 +209,20 @@ def train_classifier(
 
 ###############################################################################
 class Classifier(object):
-    def __init__(self, obj=None, lang="en_US"):
+    def __init__(self, classifier_object=None, lang="en_US", classifier_path=None):
         """
         Load the intent classifier
+
+        Parameters
+        ----------
+        obj : dict
+            Classifier object as returned by train_classifier
+        lang : str
+            Language to use (ignored if a classifier object or path is given)
+        classifier_path : str
+            Path a joblib file containing the classifier. If None, the
+            classifier will be loaded from the default location for the given
+            language.
         """
         self.tfidf_model = None
         self.classifier = None
@@ -214,47 +231,53 @@ def __init__(self, obj=None, lang="en_US"):
         if not USE_CLF:
             return
 
-        if not obj:
-            path = language.topdir(lang).joinpath("clf.joblib")
-            with path.open("rb") as file:
-                obj = joblib.load(file)
+        if not classifier_object:
+            if classifier_path is None:
+                classifier_path = language.topdir(lang).joinpath("clf.joblib")
+            with classifier_path.open("rb") as file:
+                classifier_object = joblib.load(file)
 
         cur_scipy_version = pkg_resources.get_distribution("scikit-learn").version
-        if cur_scipy_version != obj.get("scikit-learn_version"):  # pragma: no cover
+        if cur_scipy_version != classifier_object.get(
+            "scikit-learn_version"
+        ):  # pragma: no cover
             _LOGGER.warning(
                 "The classifier was built using a different scikit-learn "
                 "version (={}, !={}). The disambiguation tool could behave "
                 "unexpectedly. Consider running classifier.train_classfier()".format(
-                    obj.get("scikit-learn_version"), cur_scipy_version
+                    classifier_object.get("scikit-learn_version"), cur_scipy_version
                 )
             )
 
-        self.tfidf_model = obj["tfidf_model"]
-        self.classifier = obj["clf"]
-        self.target_names = obj["target_names"]
+        self.tfidf_model = classifier_object["tfidf_model"]
+        self.classifier = classifier_object["clf"]
+        self.target_names = classifier_object["target_names"]
 
 
 @cached
-def classifier(lang="en_US"):
+def classifier(lang: str = "en_US", classifier_path: str = None) -> Classifier:
     """
     Cached classifier object
-    :param lang:
-    :return:
+    :param lang: language
+    :param classifier_path: path to a joblib file containing the classifier
+    :return: Classifier object
     """
-    return Classifier(lang=lang)
+    return Classifier(lang=lang, classifier_path=classifier_path)
 
 
 ###############################################################################
-def disambiguate_entity(key, text, lang="en_US"):
+def disambiguate_entity(key, text, lang="en_US", classifier_path=None):
     """
     Resolve ambiguity between entities with same dimensionality.
     """
 
     new_ent = next(iter(load.entities(lang).derived[key]))
     if len(load.entities(lang).derived[key]) > 1:
-        transformed = classifier(lang).tfidf_model.transform([clean_text(text, lang)])
-        scores = classifier(lang).classifier.predict_proba(transformed).tolist()[0]
-        scores = zip(scores, classifier(lang).target_names)
+        classifier_: Classifier = classifier(lang, classifier_path)
+
+        transformed = classifier_.tfidf_model.transform([clean_text(text, lang)])
+        scores = classifier_.classifier.predict_proba(transformed).tolist()[0]
+        scores = zip(scores, classifier_.target_names)
 
         # Filter for possible names
         names = [i.name for i in load.entities(lang).derived[key]]
@@ -271,7 +294,7 @@ def disambiguate_entity(key, text, lang="en_US"):
 
 
 ###############################################################################
-def disambiguate_unit(unit, text, lang="en_US"):
+def disambiguate_unit(unit, text, lang="en_US", classifier_path=None):
     """
     Resolve ambiguity between units with same names, symbols or abbreviations.
     """
@@ -286,9 +309,11 @@ def disambiguate_unit(unit, text, lang="en_US"):
         return load.units(lang).names.get("unk")
 
     if len(new_unit) > 1:
-        transformed = classifier(lang).tfidf_model.transform([clean_text(text, lang)])
-        scores = classifier(lang).classifier.predict_proba(transformed).tolist()[0]
-        scores = zip(scores, classifier(lang).target_names)
+        classifier_: Classifier = classifier(lang, classifier_path)
+
+        transformed = classifier_.tfidf_model.transform([clean_text(text, lang)])
+        scores = classifier_.classifier.predict_proba(transformed).tolist()[0]
+        scores = zip(scores, classifier_.target_names)
 
         # Filter for possible names
         names = [i.name for i in new_unit]

quantulum3/disambiguate.py

@@ -9,13 +9,13 @@
 
 
 ###############################################################################
-def disambiguate_unit(unit_surface, text, lang="en_US"):
+def disambiguate_unit(unit_surface, text, lang="en_US", classifier_path=None):
     """
     Resolve ambiguity between units with same names, symbols or abbreviations.
     :returns (str) unit name of the resolved unit
     """
     if clf.USE_CLF:
-        base = clf.disambiguate_unit(unit_surface, text, lang).name
+        base = clf.disambiguate_unit(unit_surface, text, lang, classifier_path).name
     else:
         base = (
             load.units(lang).symbols[unit_surface]
@@ -38,13 +38,13 @@ def disambiguate_unit(unit_surface, text, lang="en_US"):
 
 
 ###############################################################################
-def disambiguate_entity(key, text, lang="en_US"):
+def disambiguate_entity(key, text, lang="en_US", classifier_path=None):
     """
     Resolve ambiguity between entities with same dimensionality.
     """
     try:
         if clf.USE_CLF:
-            ent = clf.disambiguate_entity(key, text, lang)
+            ent = clf.disambiguate_entity(key, text, lang, classifier_path)
         else:
             derived = load.entities().derived[key]
             if len(derived) > 1:

quantulum3/load.py

@@ -15,7 +15,14 @@
 TOPDIR = Path(__file__).parent or Path(".")
 
 ###############################################################################
-_CACHE_DICT = {}
+_CACHE_DICT = defaultdict(dict)
+
+
+def clear_cache():
+    """
+    Useful for testing.
+    """
+    _CACHE_DICT.clear()
 
 
 def cached(funct):
@@ -28,12 +35,15 @@ def cached(funct):
     """
     assert callable(funct)
 
-    def cached_function(lang="en_US"):
+    def cached_function(*args, **kwargs):
+        # create a hash of args and kwargs
+        args_hash = hash((args, frozenset(kwargs.items())))
+
         try:
-            return _CACHE_DICT[id(funct)][lang]
+            return _CACHE_DICT[id(funct)][args_hash]
         except KeyError:
-            result = funct(lang)
-            _CACHE_DICT[id(funct)] = {lang: result}
+            result = funct(*args, **kwargs)
+            _CACHE_DICT[id(funct)].update({args_hash: result})
             return result
 
     return cached_function

quantulum3/parser.py

@@ -227,7 +227,7 @@ def build_unit_name(dimensions, lang="en_US"):
 
 
 ###############################################################################
-def get_unit_from_dimensions(dimensions, text, lang="en_US"):
+def get_unit_from_dimensions(dimensions, text, lang="en_US", classifier_path=None):
     """
     Reconcile a unit based on its dimensionality.
     """
@@ -241,7 +241,7 @@ def get_unit_from_dimensions(dimensions, text, lang="en_US"):
         unit = cls.Unit(
             name=build_unit_name(dimensions, lang),
             dimensions=dimensions,
-            entity=get_entity_from_dimensions(dimensions, text, lang),
+            entity=get_entity_from_dimensions(dimensions, text, lang, classifier_path),
         )
 
     # Carry on original composition
@@ -268,7 +268,7 @@ def infer_name(unit):
 
 
 ###############################################################################
-def get_entity_from_dimensions(dimensions, text, lang="en_US"):
+def get_entity_from_dimensions(dimensions, text, lang="en_US", classifier_path=None):
     """
     Infer the underlying entity of a unit (e.g. "volume" for "m^3") based on
     its dimensionality.
@@ -282,7 +282,7 @@ def get_entity_from_dimensions(dimensions, text, lang="en_US"):
     final_derived = sorted(new_derived, key=lambda x: x["base"])
     key = load.get_key_from_dimensions(final_derived)
 
-    ent = dis.disambiguate_entity(key, text, lang)
+    ent = dis.disambiguate_entity(key, text, lang, classifier_path=classifier_path)
     if ent is None:
         _LOGGER.debug("\tCould not find entity for: %s", key)
         ent = cls.Entity(name="unknown", dimensions=new_derived)
@@ -299,7 +299,7 @@ def parse_unit(item, unit, slash, lang="en_US"):
 
 
 ###############################################################################
-def get_unit(item, text, lang="en_US"):
+def get_unit(item, text, lang="en_US", classifier_path=None):
     """
     Extract unit from regex hit.
     """
@@ -362,10 +362,10 @@ def get_unit(item, text, lang="en_US"):
             # Determine which unit follows
             if unit:
                 unit_surface, power = parse_unit(item, unit, slash, lang)
-                base = dis.disambiguate_unit(unit_surface, text, lang)
+                base = dis.disambiguate_unit(unit_surface, text, lang, classifier_path)
                 derived += [{"base": base, "power": power, "surface": unit_surface}]
 
-        unit = get_unit_from_dimensions(derived, text, lang)
+        unit = get_unit_from_dimensions(derived, text, lang, classifier_path)
 
     _LOGGER.debug("\tUnit: %s", unit)
     _LOGGER.debug("\tEntity: %s", unit.entity)
@@ -424,14 +424,23 @@ def is_quote_artifact(orig_text, span):
 
 ###############################################################################
 def build_quantity(
-    orig_text, text, item, values, unit, surface, span, uncert, lang="en_US"
+    orig_text,
+    text,
+    item,
+    values,
+    unit,
+    surface,
+    span,
+    uncert,
+    lang="en_US",
+    classifier_path=None,
 ):
     """
     Build a Quantity object out of extracted information.
     Takes care of caveats and common errors
     """
     return _get_parser(lang).build_quantity(
-        orig_text, text, item, values, unit, surface, span, uncert
+        orig_text, text, item, values, unit, surface, span, uncert, classifier_path
     )
 
 
@@ -479,8 +488,8 @@ def handle_consecutive_quantities(quantities, context):
         if range_span:
             if q1.unit.name == q2.unit.name or q1.unit.name == "dimensionless":
                 if (
-                    q1.uncertainty == None
-                    and q2.uncertainty == None
+                    q1.uncertainty is None
+                    and q2.uncertainty is None
                     and q1.value != q2.value
                 ):
                     a, b = (q1, q2) if q2.value > q1.value else (q2, q1)
@@ -505,9 +514,28 @@ def handle_consecutive_quantities(quantities, context):
 
 
 ###############################################################################
-def parse(text, lang="en_US", verbose=False) -> List[cls.Quantity]:
+def parse(
+    text, lang="en_US", verbose=False, classifier_path=None
+) -> List[cls.Quantity]:
     """
     Extract all quantities from unstructured text.
+
+    Parameters
+    ----------
+    text : str
+        Text to parse.
+    lang : str
+        Language of the text. Default is "en_US".
+    verbose : bool
+        If True, print debug information. Default is False.
+    classifier_path : str
+        Path to the classifier model. Default is None, which uses the default
+        model for the given language.
+
+    Returns
+    -------
+    quantities : List[Quantity]
+        List of quantities found in the text.
     """
 
     log_format = "%(asctime)s --- %(message)s"
@@ -533,10 +561,19 @@ def parse(text, lang="en_US", verbose=False) -> List[cls.Quantity]:
         try:
             uncert, values = get_values(item, lang)
 
-            unit, unit_shortening = get_unit(item, text)
+            unit, unit_shortening = get_unit(item, text, lang, classifier_path)
             surface, span = get_surface(shifts, orig_text, item, text, unit_shortening)
             objs = build_quantity(
-                orig_text, text, item, values, unit, surface, span, uncert, lang
+                orig_text,
+                text,
+                item,
+                values,
+                unit,
+                surface,
+                span,
+                uncert,
+                lang,
+                classifier_path,
             )
             if objs is not None:
                 quantities += objs