ITADN

Fix generate_mamba() argument error in interactive mode

#260Pull Requestfede-kamel 创建于 2025-11-27
F
fede-kamelcommented
## Summary Fixes #220 - TypeError when running Mamba models in interactive mode. **Root Cause:** The `interactive()` function was passing `images` as a positional argument to both `generate()` and `generate_mamba()`, but their signatures differ: - `generate(encoded_prompts, model, images, *, ...)` - accepts images as 3rd positional arg - `generate_mamba(encoded_prompts, model, *, ...)` - does NOT accept images (Mamba models don't support vision) **The Fix:** Use separate code paths for Transformer and Mamba models in `interactive()`: ```python if isinstance(model, Transformer): generated_tokens, _ = generate([tokens], model, [images], ...) else: # Mamba models don't support images generated_tokens, _ = generate_mamba([tokens], model, ...) ``` ## Changes - `src/mistral_inference/main.py`: Separate branches for Transformer vs Mamba in `interactive()` - `tests/test_main_generate_dispatch.py`: Regression tests verifying function signatures ## Test Plan - [x] Unit tests verify `generate_mamba()` doesn't accept `images` parameter - [x] Unit tests verify `generate()` accepts `images` as positional parameter - [x] Manual test with Mamba-Codestral-7B-v0.1 on NVIDIA GPU (in progress) ## Before (crash) ``` TypeError: generate_mamba() takes 2 positional arguments but 3 were given ``` ## After (works) Mamba models generate correctly in interactive mode without the images argument.
合并状态:未合并 9 条评论