A pruned and distilled variant of
allenai/Flex-math-2x7B-1T with a variable-width expert MLP. Expert 1 has been pruned from the full 11,008 intermediate size down to
5504 (50% of original width), then recovered via knowledge distillation.
For full details, see the
blog post.
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3model = AutoModelForCausalLM.from_pretrained("hbfreed/flex-math-5504", trust_remote_code=True)
4tokenizer = AutoTokenizer.from_pretrained("allenai/Flex-math-2x7B-1T")
5
6input_text = "Solve: What is 15% of 200?"
7inputs = tokenizer(input_text, return_tensors="pt")
8outputs = model.generate(**inputs, max_new_tokens=256)
9print(tokenizer.decode(outputs[0], skip_special_tokens=True))
The tokenizer is the same as the base model's.
Math-calibrated importance analysis was used — 58% of the top-2048 neurons differ between math-calibrated and general-calibrated rankings.
