| """Load the PlutoLite-v1 text adapter with its original completion format.""" |
|
|
| import argparse |
| from pathlib import Path |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--adapter", default=str(Path(__file__).resolve().parent)) |
| parser.add_argument("--base-model", default="google/gemma-4-E2B") |
| parser.add_argument("--local-files-only", action="store_true") |
| parser.add_argument("--prompt", required=True) |
| parser.add_argument("--max-new-tokens", type=int, default=256) |
| args = parser.parse_args() |
| if args.max_new_tokens < 1: |
| parser.error("--max-new-tokens must be positive") |
|
|
| import torch |
| from accelerate.hooks import remove_hook_from_module |
| from peft import PeftModel |
| from transformers import AutoTokenizer, BitsAndBytesConfig, Gemma4ForCausalLM |
|
|
| if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): |
| raise RuntimeError("This NF4 example requires a BF16-capable NVIDIA GPU.") |
| revision = "d29ff6b45f081a49ee2733a859c9c9c2d95d1a6f" |
| source_kwargs = dict( |
| revision=revision, |
| local_files_only=args.local_files_only, |
| trust_remote_code=False, |
| ) |
| tokenizer = AutoTokenizer.from_pretrained(args.base_model, **source_kwargs) |
| if tokenizer.pad_token_id is None: |
| tokenizer.pad_token = tokenizer.eos_token |
| base, loading = Gemma4ForCausalLM.from_pretrained( |
| args.base_model, |
| **source_kwargs, |
| dtype=torch.bfloat16, |
| device_map={"": "cuda:0", "model.embed_tokens_per_layer": "cpu"}, |
| attn_implementation="sdpa", |
| quantization_config=BitsAndBytesConfig( |
| load_in_4bit=True, |
| bnb_4bit_quant_type="nf4", |
| bnb_4bit_use_double_quant=True, |
| bnb_4bit_compute_dtype=torch.bfloat16, |
| llm_int8_enable_fp32_cpu_offload=True, |
| ), |
| key_mapping={r"^model\.language_model\.": "model."}, |
| output_loading_info=True, |
| ) |
| if any(loading.get(key) for key in ("missing_keys", "mismatched_keys", "error_msgs")): |
| raise RuntimeError(f"Incomplete base text model loading: {loading}") |
|
|
| |
| embedding = base.model.embed_tokens_per_layer |
| remove_hook_from_module(embedding, recurse=True) |
| embedding.to(device="cpu", dtype=torch.bfloat16) |
| embedding.register_forward_pre_hook(lambda module, inputs: (inputs[0].to("cpu"),)) |
| embedding.register_forward_hook(lambda module, inputs, output: output.to("cuda:0")) |
| placement = getattr(base, "hf_device_map", None) |
| if placement is not None: |
| del base.hf_device_map |
| try: |
| model = PeftModel.from_pretrained( |
| base, args.adapter, is_trainable=False, local_files_only=args.local_files_only |
| ) |
| finally: |
| if placement is not None: |
| base.hf_device_map = placement |
| model.eval() |
| prefix = (tokenizer.bos_token or "") + args.prompt |
| inputs = tokenizer(prefix, add_special_tokens=False, return_tensors="pt").to("cuda:0") |
| with torch.inference_mode(): |
| output = model.generate( |
| **inputs, |
| max_new_tokens=args.max_new_tokens, |
| do_sample=False, |
| pad_token_id=tokenizer.pad_token_id, |
| eos_token_id=tokenizer.eos_token_id, |
| ) |
| print(tokenizer.decode(output[0, inputs.input_ids.shape[1]:], skip_special_tokens=True)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
