minicpmv-surgery.py 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445
  1. import argparse
  2. import os
  3. import torch
  4. from transformers import AutoModel, AutoTokenizer
  5. ap = argparse.ArgumentParser()
  6. ap.add_argument("-m", "--model", help="Path to MiniCPM-V model")
  7. args = ap.parse_args()
  8. # find the model part that includes the the multimodal projector weights
  9. model = AutoModel.from_pretrained(args.model, trust_remote_code=True, local_files_only=True, torch_dtype=torch.bfloat16)
  10. checkpoint = model.state_dict()
  11. # get a list of mm tensor names
  12. mm_tensors = [k for k, v in checkpoint.items() if k.startswith("resampler")]
  13. # store these tensors in a new dictionary and torch.save them
  14. projector = {name: checkpoint[name].float() for name in mm_tensors}
  15. torch.save(projector, f"{args.model}/minicpmv.projector")
  16. clip_tensors = [k for k, v in checkpoint.items() if k.startswith("vpm")]
  17. if len(clip_tensors) > 0:
  18. clip = {name.replace("vpm.", ""): checkpoint[name].float() for name in clip_tensors}
  19. torch.save(clip, f"{args.model}/minicpmv.clip")
  20. # added tokens should be removed to be able to convert Mistral models
  21. if os.path.exists(f"{args.model}/added_tokens.json"):
  22. with open(f"{args.model}/added_tokens.json", "w") as f:
  23. f.write("{}\n")
  24. config = model.llm.config
  25. config.auto_map = {
  26. "AutoConfig": "configuration_minicpm.MiniCPMConfig",
  27. "AutoModel": "modeling_minicpm.MiniCPMModel",
  28. "AutoModelForCausalLM": "modeling_minicpm.MiniCPMForCausalLM",
  29. "AutoModelForSeq2SeqLM": "modeling_minicpm.MiniCPMForCausalLM",
  30. "AutoModelForSequenceClassification": "modeling_minicpm.MiniCPMForSequenceClassification"
  31. }
  32. model.llm.save_pretrained(f"{args.model}/model")
  33. tok = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
  34. tok.save_pretrained(f"{args.model}/model")
  35. print("Done!")
  36. print(f"Now you can convert {args.model} to a regular LLaMA GGUF file.")
  37. print(f"Also, use {args.model}/minicpmv.projector to prepare a minicpmv-encoder.gguf file.")