ml-explore / mlx-examples

Examples in the MLX framework
MIT License
5.5k stars 791 forks source link

make models/phi3.py and models/phi3small.py compatible with mypy #833

Closed wangkuiyi closed 3 weeks ago

wangkuiyi commented 3 weeks ago

Test Plan:

Run the following command:

❯ mypy \
  --explicit-package-bases \
  --ignore-missing-imports \
. | grep models/phi

Before this change, it should complain the following:

mlx_lm/models/phi3small.py:22: error: Incompatible types in assignment (expression has type "None", variable has type "int")  [assignment]
mlx_lm/models/phi3small.py:29: error: Incompatible types in assignment (expression has type "Tuple[int]", variable has type "int")  [assignment]
mlx_lm/models/phi3small.py:74: error: Name "math" is not defined  [name-defined]
mlx_lm/models/phi3small.py:176: error: "Tuple[Any, Any]" has no attribute "offset"  [attr-defined]
mlx_lm/models/phi3small.py:177: error: "Tuple[Any, Any]" has no attribute "offset"  [attr-defined]
mlx_lm/models/phi3small.py:178: error: "Tuple[Any, Any]" has no attribute "update_and_fetch"  [attr-defined]
mlx_lm/models/phi3.py:20: error: Incompatible types in assignment (expression has type "None", variable has type "int")  [assignment]
mlx_lm/models/phi3.py:68: error: Argument "short_factor" to "SuScaledRotaryEmbedding" has incompatible type "Union[float, str]"; expected "Union[List[float], float]"  [arg-type]
mlx_lm/models/phi3.py:69: error: Argument "long_factor" to "SuScaledRotaryEmbedding" has incompatible type "Union[float, str]"; expected "Union[List[float], float]"  [arg-type]
mlx_lm/models/phi3.py:73: error: Unsupported operand types for / ("int" and "str")  [operator]
mlx_lm/models/phi3.py:73: note: Right operand is of type "Union[float, str]"
mlx_lm/models/phi3.py:101: error: "Tuple[Any, Any]" has no attribute "offset"  [attr-defined]
mlx_lm/models/phi3.py:102: error: "Tuple[Any, Any]" has no attribute "offset"  [attr-defined]
mlx_lm/models/phi3.py:103: error: "Tuple[Any, Any]" has no attribute "update_and_fetch"  [attr-defined]

After this change, the above command should complain no more.