ml-explore / mlx-examples

Examples in the MLX framework
MIT License
5.92k stars 838 forks source link

Fix mypy errors with models/{qwen2,qwen2_moe,startcoder2}.py #835

Closed wangkuiyi closed 3 months ago

wangkuiyi commented 3 months ago

Before this change, we have

❯ mypy \
  --explicit-package-bases \
  --ignore-missing-imports \
. | grep 'qwen2'
llms/mlx_lm/models/qwen2_moe.py:25: error: Incompatible types in assignment (expression has type "None", variable has type "int")  [assignment]
llms/mlx_lm/models/qwen2_moe.py:82: error: "Tuple[Any, Any]" has no attribute "offset"  [attr-defined]
llms/mlx_lm/models/qwen2_moe.py:83: error: "Tuple[Any, Any]" has no attribute "offset"  [attr-defined]
llms/mlx_lm/models/qwen2_moe.py:84: error: "Tuple[Any, Any]" has no attribute "update_and_fetch"  [attr-defined]
llms/mlx_lm/models/qwen2.py:19: error: Incompatible types in assignment (expression has type "None", variable has type "int")  [assignment]
llms/mlx_lm/models/qwen2.py:55: error: Unsupported operand types for / ("int" and "str")  [operator]
llms/mlx_lm/models/qwen2.py:55: note: Right operand is of type "Union[float, str]"
llms/mlx_lm/models/qwen2.py:82: error: "Tuple[Any, Any]" has no attribute "offset"  [attr-defined]
llms/mlx_lm/models/qwen2.py:83: error: "Tuple[Any, Any]" has no attribute "offset"  [attr-defined]
llms/mlx_lm/models/qwen2.py:84: error: "Tuple[Any, Any]" has no attribute "update_and_fetch"  [attr-defined]

and

❯ mypy \
  --explicit-package-bases \
  --ignore-missing-imports \
. | grep 'starcoder2'
llms/mlx_lm/models/starcoder2.py:58: error: "Tuple[Any, Any]" has no attribute "offset"  [attr-defined]
llms/mlx_lm/models/starcoder2.py:59: error: "Tuple[Any, Any]" has no attribute "offset"  [attr-defined]
llms/mlx_lm/models/starcoder2.py:60: error: "Tuple[Any, Any]" has no attribute "update_and_fetch"  [attr-defined]

After this change, the above errors disappear.

TODO

Consider change the following dictionary typed field rope_scaling into a nested dataclass in a future PR.

@dataclass
class ModelArgs(BaseModelArgs):
    rope_scaling: Optional[Dict[str, Union[float, str]]] = None