from mlx_lm import load as mlx_lm_load
from dflash.model_mlx import load_draft, stream_generate
from mlx_lm.sample_utils import make_sampler

target, tokenizer = mlx_lm_load("speakleash/Bielik-11B-v3.0-Instruct-MLX-8bit")
draft = load_draft("gangel/Bielik-11B-v3.0-DFlash-MLX-8bit")

prompt = tokenizer.apply_chat_template(
    [{"role": "user", "content": "Jaka jest stolica Polski?"}],
    tokenize=False, add_generation_prompt=True,
)

for r in stream_generate(target, draft, tokenizer, prompt,
                          block_size=16, max_tokens=256, temperature=0.0):
    print(r.text, end="", flush=True)