Views
No views yet
1from transformers import AutoConfig, AutoModel, logging
2from transformers import AutoModel, AutoTokenizer
3import torch
4from PIL import Image
5import os
6
7logging.set_verbosity_error() # silence HF info spam
8
9MODEL_ID = "openbmb/MiniCPM-o-2_6"
10
11device = "cpu"
12cfg = AutoConfig.from_pretrained(MODEL_ID, trust_remote_code=True)
13
14cfg.hidden_size = 24 * 6
15#cfg.hidden_size = 128
16cfg.num_heads = 1
17cfg.num_hidden_layers = 28
18cfg.intermediate_size = 16
19cfg.num_attention_heads=24
20
21cfg.vision_config.hidden_size = 8
22cfg.vision_config.num_hidden_layers = 1
23cfg.vision_config.num_attention_heads = 1
24cfg.vision_config.intermediate_size = 8
25#cfg.vision_config.image_size = 100
26
27cfg.audio_config.encoder_layers = 1
28cfg.audio_config.decoder_layers = 1
29cfg.audio_config.decoder_ffn_dim = 1024
30#cfg.audio_config.d_model = 32
31#cfg.audio_config.encoder_ffn_dim = 1024
32#cfg.audio_config.use_bfloat16=True
33
34cfg.tts_config.llm_dim = 16
35cfg.tts_config.hidden_size = 12
36
37cfg.tts_config.llm_dim = 4 # keep small (interface with LM)
38cfg.tts_config.hidden_size = 8 # shrink internal TTS width
39cfg.tts_config.intermediate_size = 4 # shrink FFN
40cfg.tts_config.num_layers = 1 # minimum, keeps a single block
41cfg.tts_config.num_heads = 1 # avoid multi-head blowup
42cfg.tts_config.num_hidden_layers = 1
43cfg.tts_config.num_mel_bins = 10
44cfg.tts_config.num_attention_heads = 1
45cfg.tts_config.num_text_tokens = 20
46cfg.tts_config.num_audio_tokens = 10
47#cfg.tts_config.use_bfloat16=True
48
49model = AutoModel.from_config(cfg, trust_remote_code=True)
50
51# cast to bfloat16
52model = model.to(dtype=torch.bfloat16, device=device)
53
54print("Built tiny MiniCPM-o model on", device)
55
56tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
57
58output_dir = "./tiny-random-minicpmo-new-version"
59os.makedirs(output_dir, exist_ok=True)
60model.save_pretrained(output_dir, safe_serialization=True)
61tokenizer.save_pretrained(output_dir)
62model.processor.save_pretrained(output_dir)