fix detokenizer extra tokens

Signed-off-by: Alex Chi Z <[email protected]>

Author: skyzh

main.py

@@ -13,34 +13,43 @@
 parser.add_argument("--device", type=str, default="gpu")
 args = parser.parse_args()
 
+use_mlx = False
 if args.solution == "tiny_llm":
     from tiny_llm import Qwen2Model, simple_generate
-
     print("Using your tiny_llm solution")
 elif args.solution == "tiny_llm_week1_ref" or args.solution == "week1_ref":
     from tiny_llm_week1_ref import Qwen2Model, simple_generate
-
     print("Using tiny_llm_week1_ref solution")
 elif args.solution == "tiny_llm_week2_ref" or args.solution == "week2_ref":
     from tiny_llm_week2_ref import Qwen2Model, simple_generate
-
     print("Using tiny_llm_week2_ref solution")
+elif args.solution == "mlx":
+    use_mlx = True
+    from mlx_lm.generate import stream_generate
+    print("Using the original mlx model")
 else:
     raise ValueError(f"Solution {args.solution} not supported")
 
 mlx_model, tokenizer = load(
     args.model,
     tokenizer_config={"eos_token": "<|im_end|>"},
-    model_config={"tie_word_embeddings": False, "rope_traditional": True},
+    model_config={"tie_word_embeddings": False, "rope_traditional": False},
 )
 
 with mx.stream(mx.gpu if args.device == "gpu" else mx.cpu):
-    tiny_llm_model = Qwen2Model(mlx_model)
+    if use_mlx:
+        tiny_llm_model = mlx_model
+    else:
+        tiny_llm_model = Qwen2Model(mlx_model)
     messages = [
         {"role": "system", "content": "You are a helpful assistant."},
         {"role": "user", "content": args.prompt},
     ]
     prompt = tokenizer.apply_chat_template(
         messages, tokenize=False, add_generation_prompt=True
     )
-    simple_generate(tiny_llm_model, tokenizer, prompt)
+    if not use_mlx:
+        simple_generate(tiny_llm_model, tokenizer, prompt)
+    else:
+        for resp in stream_generate(tiny_llm_model, tokenizer, prompt):
+            print(resp.text, end="", flush=True)

src/tiny_llm_week2_ref/generate.py

@@ -23,9 +23,10 @@ def _step(model, y, offset):
     # generate/decode
     while True:
         token, _ = _step(model, tokens, offset)
-        offset += tokens.size
-        tokens = token
+        if offset != 0:
+            detokenizer.add_token(token.item())
+            print(detokenizer.last_segment, end="", flush=True)
         if token.item() == tokenizer.eos_token_id:
             break
-        detokenizer.add_token(token.item())
-        print(detokenizer.last_segment, end="", flush=True)
+        offset += tokens.size
+        tokens = token