File size: 4,233 Bytes
259c99a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
from safetensors import safe_open
from safetensors.torch import save_file
from huggingface_hub import hf_hub_download, login
import os
from tqdm import tqdm
def load_gemini_weights(repo_id):
print("Downloading Gemini Nano weights...")
gemini_model_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors")
return gemini_model_path
def adapt_gemini_to_gemma(gemini_path, output_path, custom_config):
print("Adapting Gemini weights to Gemma format...")
with safe_open(gemini_path, framework="pt", device="cpu") as f:
gemini_keys = list(f.keys())
# Process embedding layer
embed_weight = f.get_tensor('model.embed_tokens.weight')
vocab_size = embed_weight.size(0) // 2048
embed_weight = embed_weight[:vocab_size * 2048].view(vocab_size, 2048)
adapted_weights = {'model.embed_tokens.weight': embed_weight,
'lm_head.weight': embed_weight.clone()}
# Process other layers
for key in tqdm(gemini_keys, desc="Processing layers"):
if key.startswith('model.layers.'):
parts = key.split('.')
layer_num = int(parts[2])
if layer_num >= custom_config.num_hidden_layers:
continue
weight = f.get_tensor(key)
adapted_weights[key] = weight
elif key == 'model.norm.weight':
adapted_weights[key] = f.get_tensor(key)
# Save adapted weights with metadata
print("Saving adapted weights...")
metadata = {"format": "pt"}
save_file(adapted_weights, output_path, metadata=metadata)
return adapted_weights
def create_custom_gemma_config(gemma_repo, gemini_path):
with safe_open(gemini_path, framework="pt", device="cpu") as f:
embed_weight = f.get_tensor('model.embed_tokens.weight')
ffn_weight = f.get_tensor('model.layers.0.mlp.gate_proj.weight')
custom_config = AutoConfig.from_pretrained(gemma_repo)
custom_config.vocab_size = embed_weight.size(0) // 2048
custom_config.intermediate_size = ffn_weight.size(0)
custom_config.num_hidden_layers = 32 # Assuming Gemini Nano has 32 layers
return custom_config
def test_model(model_path, tokenizer, prompt, max_length=50):
print("Testing the adapted model...")
# Load the model in 8-bit quantization
model = AutoModelForCausalLM.from_pretrained(
model_path,
device_map="auto",
load_in_8bit=True,
ignore_mismatched_sizes=True
)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
outputs = model.generate(**inputs, max_length=max_length)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
return generated_text
def main():
gemini_repo = "QuietImpostor/Gemini-Nano-Safetensors"
gemma_repo = "google/gemma-2b-it"
output_path = "/kaggle/temp/gemini-gemmafied/"
login(token="hf_...")
# Load Gemini Nano weights
gemini_path = load_gemini_weights(gemini_repo)
# Create custom Gemma config
custom_config = create_custom_gemma_config(gemma_repo, gemini_path)
# Adapt Gemini weights to Gemma format
adapted_weights_path = os.path.join(output_path, "model.safetensors")
os.makedirs(output_path, exist_ok=True)
adapt_gemini_to_gemma(gemini_path, adapted_weights_path, custom_config)
# Save the custom config
custom_config.save_pretrained(output_path)
# Load Gemini Nano tokenizer and save it
print("Saving tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(gemini_repo)
tokenizer.save_pretrained(output_path)
print("Adaptation complete!")
# Test the model
prompt = "The future of artificial intelligence is"
generated_text = test_model(output_path, tokenizer, prompt)
print(f"Prompt: {prompt}")
print(f"Generated text: {generated_text}")
if __name__ == "__main__":
main() |