feat: function signature matches functionalities

sileod · Sep 19, 2023 · 89aadf4 · 89aadf4
1 parent a253a11
commit 89aadf4
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/src/tasknet/utils.py b/src/tasknet/utils.py
@@ -152,7 +152,14 @@ def search_module(m, name, mode="attr", lowercase=True):
         raise ValueError('mode must be "attr" or "class"')
 
 
-def load_pipeline(model_name, task_name, adapt_task_embedding=True, multilingual=False):
+def load_pipeline(
+    model_name: str,
+    task_name: str,
+    adapt_task_embedding: bool = True,
+    multilingual: bool = False,
+    device: int = -1,
+    return_all_scores: bool = False,
+) -> TextClassificationPipeline:
     """Load Text Classification Pipeline for a Specified Model.
 
     Load a text classification pipeline for the specified model and task. If