From 1b8869e3c2b20a910523e8e48d0a39c5a4f9e7da 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..7d8afba5 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 = lit_model.Model + + 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)