Upload conversion script

Adding the script used to convert Gemini Nano to Gemma weights as requested by ethanc8.

gemmafy_gemini.py

@@ -0,0 +1,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()