From 0e354cbde72e1e1d7beb7065bd8d298b88359840 Mon Sep 17 00:00:00 2001 From: Fan Ye Date: Wed, 25 Sep 2024 12:31:16 -0700 Subject: [PATCH] LIT:Fix type hints for model initializers. PiperOrigin-RevId: 678805235 --- lit_nlp/app.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/lit_nlp/app.py b/lit_nlp/app.py index 94343fe9..7fe822df 100644 --- a/lit_nlp/app.py +++ b/lit_nlp/app.py @@ -17,6 +17,7 @@ import collections from collections.abc import Callable, Iterable, Mapping, Sequence import functools +import inspect import math import os import random @@ -520,7 +521,10 @@ def _create_model( raise_for_unsupported=True, ) - return_type = get_type_hints(model_initializer)['return'] + return_type = dict[str, Any] + + if inspect.isfunction(model_initializer): + return_type = get_type_hints(model_initializer)['return'] if Mapping in return_type.__mro__: model_initializer = cast(MultipleModelLoader, model_initializer)