zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-07-23 20:51:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

don't drop last tokens

Browse Source
This commit is contained in:
Awni Hannun 2023-12-14 08:27:44 -08:00
parent 1613e608a9
commit 840c0c36c2
1 changed files with 5 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
phi2/model.py
View File

@ -159,11 +159,8 @@ def generate(prompt: mx.array, model: Phi2, temp: Optional[float] = 0.0):
def load_model(): def load_model():
model = Phi2(ModelArgs()) model = Phi2(ModelArgs())
weights = mx.load("weights.npz") weights = mx.load("weights.npz")
weights = tree_unflatten(list(weights.items())) model.update(tree_unflatten(list(weights.items())))
model.update(weights)
tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True) tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
return model, tokenizer return model, tokenizer
@ -215,3 +212,7 @@ if __name__ == "__main__":
s = tokenizer.decode([t.item() for t in tokens]) s = tokenizer.decode([t.item() for t in tokens])
print(s, end="", flush=True) print(s, end="", flush=True)
tokens = [] tokens = []
mx.eval(tokens)
s = tokenizer.decode([t.item() for t in tokens])
print(s, flush=True)