Views
No views yet
1import torch
2import gc
3import copy
4import os
5from google.colab import drive
6from huggingface_hub import login
7
8from transformers import (
9 AutoModel,
10 AutoConfig,
11 AutoModelForCausalLM,
12 AutoTokenizer
13)
14
15drive.mount('/content/drive')
16
17# Configuration
18SOURCE_MODEL_ID = "google/gemma-3-4b-it"
19OUTPUT_DIR = "/content/drive/MyDrive/gemma-3-4b-it-novision"
20
21def convert_gemma3_multimodal_to_text():
22 print(f"Loading config from {SOURCE_MODEL_ID}...")
23
24 # 1. Prepare the Configuration
25 full_config = AutoConfig.from_pretrained(SOURCE_MODEL_ID, trust_remote_code=True)
26
27 if not hasattr(full_config, 'text_config'):
28 raise ValueError("Config does not contain 'text_config'. Is this a Gemma 3 Multimodal model?")
29
30 # Deepcopy the text_config
31 new_config = copy.deepcopy(full_config.text_config)
32
33 # UPDATE: Gemma 3 text backbone is compatible with Gemma 2 architecture
34 # There is no "Gemma3ForCausalLM" class in transformers yet.
35 new_config.architectures = ["Gemma2ForCausalLM"]
36 # Ensure model_type is set correctly for the class mapping
37 new_config.model_type = "gemma2"
38
39 print("Configuration prepared.")
40
41 # 2. Load the Original Model
42 print(f"Loading original model weights (CPU)...")
43 multimodal_model = AutoModel.from_pretrained(
44 SOURCE_MODEL_ID,
45 torch_dtype=torch.bfloat16,
46 device_map="cpu",
47 trust_remote_code=True
48 )
49
50 # 3. Create the State Dictionary Mapping
51 print("Extracting Language Model weights...")
52 full_sd = multimodal_model.state_dict()
53 text_sd = {}
54
55 prefix_to_remove = "language_model."
56 embed_weight = None
57
58 keys_dropped = 0
59 keys_kept = 0
60
61 for key, value in full_sd.items():
62 # Drop vision tower weights
63 if "vision_tower" in key:
64 keys_dropped += 1
65 continue
66
67 # Map language model weights
68 if key.startswith(prefix_to_remove):
69 # Strip "language_model."
70 stripped_key = key[len(prefix_to_remove):]
71
72 # Identify embeddings for weight tying
73 if "embed_tokens.weight" in stripped_key:
74 embed_weight = value
75
76 # Rename logic:
77 if stripped_key.startswith("lm_head"):
78 new_key = stripped_key
79 else:
80 new_key = f"model.{stripped_key}"
81
82 text_sd[new_key] = value
83 keys_kept += 1
84 else:
85 keys_dropped += 1
86
87 # 4. Handle Weight Tying
88 if "lm_head.weight" not in text_sd:
89 print("Notice: 'lm_head.weight' not found. Creating it from embeddings (Weight Tying).")
90 if embed_weight is not None:
91 text_sd["lm_head.weight"] = embed_weight
92 else:
93 raise ValueError("Could not find embedding weights to tie to lm_head!")
94
95 print(f"Extraction complete. Kept {keys_kept} keys. Dropped {keys_dropped} keys.")
96
97 # 5. Clean up Memory
98 del multimodal_model
99 del full_sd
100 gc.collect()
101
102 # 6. Create the Text-Only Model
103 print("Instantiating new Text-Only model...")
104 # We load as Gemma2ForCausalLM
105 text_model = AutoModelForCausalLM.from_config(new_config)
106
107 # Load the filtered weights
108 print("Loading extracted weights into new architecture...")
109 text_model.load_state_dict(text_sd, strict=True)
110
111 text_model.to(dtype=torch.bfloat16)
112 text_model.eval()
113
114 # 7. Save Model and Tokenizer
115 print(f"Saving model to {OUTPUT_DIR}...")
116 text_model.save_pretrained(OUTPUT_DIR)
117
118 print("Saving tokenizer...")
119 tokenizer = AutoTokenizer.from_pretrained(SOURCE_MODEL_ID)
120 tokenizer.save_pretrained(OUTPUT_DIR)
121
122 print(f"SUCCESS! Text-only model saved to: {OUTPUT_DIR}")
123
124# Run the conversion
125convert_gemma3_multimodal_to_text()