zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-29 18:26:37 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add hf generation for comparison

Browse Source
This commit is contained in:
Juarez Bochi 2023-12-18 11:35:16 -05:00
parent 305a52dde8
commit 5ae339f6d2
No known key found for this signature in database
GPG Key ID: 34CCBB77DC8BEBB6
1 changed files with 21 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
t5/hf_t5.py
View File

@ -1,9 +1,9 @@
from transformers import T5EncoderModel, AutoTokenizer
from transformers import T5ForConditionalGeneration, T5EncoderModel, AutoTokenizer
import argparse
def run(t5_model: str):
def embed(t5_model: str):
batch = [
"translate English to German: That is good.",
"This is an example of T5 working on MLX.",
@ -22,15 +22,33 @@ def run(t5_model: str):
print()
def generate(t5_model: str):
prompt = "translate English to German: As much as six inches of rain could fall in the New York City region through Monday morning, and officials warned of flooding along the coast."
tokenizer = AutoTokenizer.from_pretrained(t5_model)
torch_model = T5ForConditionalGeneration.from_pretrained(t5_model)
torch_tokens = tokenizer(prompt, return_tensors="pt", padding=True).input_ids
outputs = torch_model.generate(torch_tokens)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Run the T5 model using Hugging Face Transformers."
)
parser.add_argument(
"--encode-only",
action="store_true",
help="Only run the encoder and print the embeddings.",
default=False,
)
parser.add_argument(
"--model",
default="t5-small",
help="The huggingface name of the T5 model to save.",
)
args = parser.parse_args()
if args.encode_only:
embed(args.model)
else:
generate(args.model)
run(args.model)