Code in the first example for Python SDK do not pass type check on load()'s returned values

[rija]

Hi, the example run fine, but when I want to use same pattern in my own programs which I always type check, errors are thrown.

    model, tokenizer = load("mlx-community/Mistral-7B-Instruct-v0.3-4bit") 
    prompt = "Write a short story about ..."

    messages = [{"role": "user", "content": prompt}]
    prompt = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
    )

    text = generate(model, tokenizer, prompt, verbose=True)
    print(text)

With pyrefly, the following error is thrown:

$ hatch check types
 INFO Checking project configured at `/Users/xxxxx/Library/Application Support/hatch/env/virtual/.config/asA_QThn/pyrefly.toml`
ERROR Cannot unpack tuple[Module, TokenizerWrapper] | tuple[Module, TokenizerWrapper, dict[str, Any]] (of size 3) into 2 values [bad-unpacking]
 --> src/myprog/myprog:6:5
  |
6 |     model, tokenizer = load("mlx-community/Mistral-7B-Instruct-v0.3-4bit")
  |     ^^^^^^^^^^^^^^^^

With mypy --strict, the following error is thrown:

$ hatch check types
src/myprog/myprog:6: error: Too many values to unpack (2 expected, 3 provided)  [misc]

My current workaround is to disable type check for that line (# type: ignore), but it's not ideal, as LLM inference is critical path.

My understanding, is that which one of the two components of the union is returned depends on whether return_config parameter is True or False.

On Stack Overflow¹, the accepted solution mentions using overload to fix the problem.

Something like that I guess:

@overload
def load(return_config: Literal[True]) -> tuple[Module, TokenizerWrapper, dict[str, Any]]: 
    ...
@overload
def load(return_config: Literal[False] = ...) -> tuple[Module, TokenizerWrapper]:
    ...
    
def load(return_config: bool = False) -> tuple[Module, TokenizerWrapper] | tuple[Module, TokenizerWrapper, dict[str, Any]]:
    ...
       if return_config:
        return model, tokenizer, config
    else:
        return model, tokenizer

Footnotes

1. https://stackoverflow.com/a/72415876/6518111 ↩