Views
No views yet
1# Download the chat script from this repository
2wget https://huggingface.co/MistyozAI/CosmicFish-90M/resolve/main/chat.py
3
4# Install dependencies
5pip install transformers huggingface-hub termcolor safetensors
6
7# Run the chat interface (automatically downloads model)
8python chat.pychat.py script handles all model loading, generation, and provides the best chat experience with live streaming, repetition penalty, and conversation commands.pip install transformers huggingface-hub termcolor safetensors torch1from transformers import GPT2Tokenizer
2from huggingface_hub import snapshot_download
3from safetensors.torch import load_file
4import torch
5import json
6import os
7
8# Download model from Hugging Face Hub
9cache_dir = snapshot_download(repo_id="MistyozAI/CosmicFish-90M")
10
11# Load tokenizer
12tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
13
14# Load config
15with open(os.path.join(cache_dir, "config.json")) as f:
16 config_dict = json.load(f)
17
18# Load model weights from safetensors
19state_dict = load_file(os.path.join(cache_dir, "model.safetensors"))
20
21# Note: Full model class available in the repository
22print("Model downloaded and ready for use!")1def generate_with_repetition_penalty(model, tokenizer, prompt, max_tokens=100, temperature=0.5, penalty=1.2):
2 input_ids = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0)
3 generated = input_ids.clone()
4
5 for _ in range(max_tokens):
6 with torch.no_grad():
7 logits, _ = model(generated)
8
9 next_token_logits = logits[:, -1, :] / temperature
10
11 # Apply repetition penalty
12 if penalty > 1.0:
13 for token_id in set(generated[0].tolist()):
14 if next_token_logits[0, token_id] > 0:
15 next_token_logits[0, token_id] /= penalty
16 else:
17 next_token_logits[0, token_id] *= penalty
18
19 probs = torch.nn.functional.softmax(next_token_logits, dim=-1)
20 next_token = torch.multinomial(probs, num_samples=1)
21
22 if next_token.item() == tokenizer.eos_token_id:
23 break
24
25 generated = torch.cat([generated, next_token], dim=1)
26
27 return tokenizer.decode(generated[0], skip_special_tokens=True)1from safetensors.torch import load_file
2from modeling_cosmicfish import CosmicFish, CosmicConfig
3import json
4
5def load_cosmicfish_model(model_path):
6 # Load config
7 with open(os.path.join(model_path, "config.json")) as f:
8 config_dict = json.load(f)
9
10 # Create model config
11 config = CosmicConfig(
12 vocab_size=config_dict["vocab_size"],
13 block_size=config_dict["block_size"],
14 n_layer=config_dict["n_layer"],
15 n_head=config_dict["n_head"],
16 n_embd=config_dict["n_embd"],
17 bias=config_dict["bias"],
18 dropout=0.0,
19 use_rotary=config_dict["use_rotary"],
20 use_swiglu=config_dict["use_swiglu"],
21 use_gqa=config_dict["use_gqa"],
22 n_query_groups=config_dict["n_query_groups"]
23 )
24
25 # Create model
26 model = CosmicFish(config)
27
28 # Load weights from safetensors (secure format)
29 state_dict = load_file(os.path.join(model_path, "model.safetensors"))
30
31 # Handle weight sharing (lm_head.weight shares with transformer.wte.weight)
32 if 'lm_head.weight' not in state_dict and 'transformer.wte.weight' in state_dict:
33 state_dict['lm_head.weight'] = state_dict['transformer.wte.weight']
34
35 model.load_state_dict(state_dict)
36 model.eval()
37
38 return model1def chat_with_model():
2 conversation = []
3
4 while True:
5 user_input = input("You: ")
6 if user_input.lower() in ['quit', 'exit']:
7 break
8
9 context = "Below is a conversation between a human and an AI assistant.\n\n"
10 for human, ai in conversation:
11 context += f"Human: {human}\nAssistant: {ai}\n\n"
12 context += f"Human: {user_input}\nAssistant:"
13
14 # Generate response with repetition penalty
15 response = generate_with_repetition_penalty(
16 model, tokenizer, context,
17 max_tokens=150, temperature=0.7, penalty=1.2
18 )
19
20 # Extract just the assistant's response
21 response = response.split("Assistant:")[-1].split('\n')[0].strip()
22 print(f"CosmicFish: {response}")
23
24 conversation.append((user_input, response))
25
26chat_with_model()