Updated for safetensors

README.md

@@ -32,7 +32,7 @@ A 120M parameter language model with modern architecture improvements developed
 wget https://huggingface.co/MistyozAI/CosmicFish-120M/resolve/main/chat.py
 
 # Install dependencies
-pip install transformers huggingface-hub termcolor
+pip install transformers huggingface-hub termcolor safetensors
 
 # Run the chat interface (automatically downloads model)
 python chat.py
@@ -49,13 +49,14 @@ The `chat.py` script handles all model loading, generation, and provides the bes
 - **Training Data**: CosmicSet 1.0
 - **Developer**: Mistyoz AI
 - **Repository**: MistyozAI/CosmicFish-120M
+- **Format**: Safetensors
 
 ## Usage
 
 ### Installation
 
 ```bash
-pip install transformers huggingface-hub termcolor
+pip install transformers huggingface-hub termcolor safetensors
 ```
 
 ### Quick Chat Interface
@@ -63,6 +64,7 @@ pip install transformers huggingface-hub termcolor
 ```python
 from transformers import GPT2Tokenizer
 from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
 import torch
 import json
 import os
@@ -77,8 +79,8 @@ tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 with open(os.path.join(cache_dir, "config.json")) as f:
     config_dict = json.load(f)
 
-# Load model weights
-state_dict = torch.load(os.path.join(cache_dir, "pytorch_model.bin"), map_location="cpu")
+# Load model weights from safetensors
+state_dict = load_file(os.path.join(cache_dir, "model.safetensors"))
 
 # Note: Full model class available in the repository
 print("Model downloaded and ready for use!")
@@ -87,7 +89,7 @@ print("Model downloaded and ready for use!")
 ### Advanced Generation with Repetition Penalty
 
 ```python
-def generate_with_repetition_penalty(model, tokenizer, prompt, max_tokens=100, temperature=0.7, penalty=1.2):
+def generate_with_repetition_penalty(model, tokenizer, prompt, max_tokens=100, temperature=0.5, penalty=1.2):
     input_ids = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0)
     generated = input_ids.clone()
     
@@ -116,6 +118,49 @@ def generate_with_repetition_penalty(model, tokenizer, prompt, max_tokens=100, t
     return tokenizer.decode(generated[0], skip_special_tokens=True)
 ```
 
+### Loading Model with Safetensors
+
+```python
+from safetensors.torch import load_file
+from modeling_cosmicfish import CosmicFish, CosmicConfig
+import json
+
+def load_cosmicfish_model(model_path):
+    # Load config
+    with open(os.path.join(model_path, "config.json")) as f:
+        config_dict = json.load(f)
+    
+    # Create model config
+    config = CosmicConfig(
+        vocab_size=config_dict["vocab_size"],
+        block_size=config_dict["block_size"], 
+        n_layer=config_dict["n_layer"],
+        n_head=config_dict["n_head"],
+        n_embd=config_dict["n_embd"],
+        bias=config_dict["bias"],
+        dropout=0.0,
+        use_rotary=config_dict["use_rotary"],
+        use_swiglu=config_dict["use_swiglu"],
+        use_gqa=config_dict["use_gqa"],
+        n_query_groups=config_dict["n_query_groups"]
+    )
+    
+    # Create model
+    model = CosmicFish(config)
+    
+    # Load weights from safetensors (secure format)
+    state_dict = load_file(os.path.join(model_path, "model.safetensors"))
+    
+    # Handle weight sharing (lm_head.weight shares with transformer.wte.weight)
+    if 'lm_head.weight' not in state_dict and 'transformer.wte.weight' in state_dict:
+        state_dict['lm_head.weight'] = state_dict['transformer.wte.weight']
+    
+    model.load_state_dict(state_dict)
+    model.eval()
+    
+    return model
+```
+
 ### Chat Interface
 
 ```python
@@ -166,13 +211,23 @@ CosmicFish uses several modern improvements over standard transformers:
 ## Performance
 
 - **Speed**: Varies by hardware (not benchmarked)
-- **Memory**: ~500MB RAM (FP16)
+- **Memory**: ~500MB RAM
 - **File Size**: 243MB
+- **Loading**: Fast and secure with safetensors
+
+## Model Format
+
+This model uses **safetensors** format for:
+- **Security**: Safe loading without arbitrary code execution
+- **Performance**: Faster loading compared to pickle-based formats
+- **Memory efficiency**: Zero-copy loading when possible
+- **Cross-platform compatibility**: Works consistently across different environments
 
 ## Limitations
 
 - Small model size (120M parameters) may produce less accurate responses
 - 512 token context limit
+- English only
 - Training data cutoff applies
 - May generate incorrect information
 - Cannot browse internet or access real-time data