Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

To make the hardcoded tokenizer path configurable via a parameter #57

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

To make the hardcoded tokenizer path configurable via a parameter #57

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8dffaaa
Update benchmark_generation_mamba_simple.py
A-transformer Nov 19, 2024
a23fd96
Update benchmark_generation_mamba_simple.py
A-transformer Nov 19, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions mamba/benchmarks/benchmark_generation_mamba_simple.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

parser = argparse.ArgumentParser(description="Generation benchmarking")
parser.add_argument("--model-name", type=str, default="state-spaces/mamba-130m")
parser.add_argument("--tokenizer-path", type=str, default=None, help="Path to the tokenizer")
parser.add_argument("--prompt", type=str, default=None)
parser.add_argument("--promptlen", type=int, default=100)
parser.add_argument("--genlen", type=int, default=100)
Expand All @@ -33,11 +34,14 @@
is_mamba = args.model_name.startswith("state-spaces/mamba-") or "mamba" in args.model_name

if is_mamba:
tokenizer = AutoTokenizer.from_pretrained("/home/zhulianghui/VisionProjects/mamba/ckpts/gpt-neox-20b-tokenizer")
if args.tokenizer_path is None:
raise ValueError("For Mamba models, --tokenizer-path must be specified.")
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
model = MambaLMHeadModel.from_pretrained(args.model_name, device=device, dtype=dtype)
else:
tokenizer = AutoTokenizer.from_pretrained(args.model_name)
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path or args.model_name)
model = AutoModelForCausalLM.from_pretrained(args.model_name, device_map={"": device}, torch_dtype=dtype)

model.eval()
print(f"Number of parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")

Expand Down