use woodwork typing for elmo primitive i/o

alteryx · Sep 30, 2021 · 755d10b · 755d10b
1 parent 2537a88
commit 755d10b
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/nlp_primitives/elmo.py b/nlp_primitives/elmo.py
@@ -1,5 +1,6 @@
 from featuretools.primitives.base import TransformPrimitive
-from featuretools.variable_types import NaturalLanguage, Numeric
+from woodwork.column_schema import ColumnSchema
+from woodwork.logical_types import Double, NaturalLanguage
 
 
 class Elmo(TransformPrimitive):
@@ -25,8 +26,8 @@ class Elmo(TransformPrimitive):
         [-0.3457, -0.4546, 0.2538]
     """
     name = "elmo"
-    input_types = [NaturalLanguage]
-    return_type = Numeric
+    input_types = [ColumnSchema(logical_type=NaturalLanguage)]
+    return_type = ColumnSchema(logical_type=Double, semantic_tags="numeric")
 
     def __init__(self):
         self.handle = "https://tfhub.dev/google/elmo/2"