Quantization made by Richard Erkhov.
This model is for debugging. It is randomly initialized using the config from
tiiuae/falcon-mamba-7b but with smaller size.
1import os
2
3import torch
4
5from huggingface_hub import create_repo, upload_folder
6from transformers import (
7 AutoModelForCausalLM,
8 AutoTokenizer,
9 GenerationConfig,
10 AutoConfig,
11 pipeline,
12 set_seed,
13)
14
15model_id = "tiiuae/falcon-mamba-7b"
16repo_id = "yujiepan/falcon-mamba-tiny-random"
17save_path = f"/tmp/{repo_id}"
18os.system(f'rm -rf {save_path}')
19
20config = AutoConfig.from_pretrained(model_id)
21config.use_cache = True
22config.num_hidden_layers = 2
23config.hidden_size = 8
24config.intermediate_size = 16
25config.state_size = 8
26
27tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
28tokenizer.save_pretrained(save_path)
29
30model = AutoModelForCausalLM.from_config(
31 config, torch_dtype=torch.bfloat16,
32 trust_remote_code=True,
33)
34model.generation_config = GenerationConfig.from_pretrained(
35 model_id,
36 trust_remote_code=True,
37)
38
39set_seed(42)
40num_params = 0
41with torch.no_grad():
42 for name, p in sorted(model.named_parameters()):
43 print(name, p.shape)
44 torch.nn.init.uniform_(p, -0.5, 0.5)
45 num_params += p.numel()
46print("Total number of parameters:", num_params)
47model.save_pretrained(save_path)
48
49pipe = pipeline(
50 "text-generation",
51 model=save_path,
52 device="cpu",
53 trust_remote_code=True,
54 max_new_tokens=20,
55)
56print(pipe("Hello World!"))
57
58# create_repo(repo_id, exist_ok=True)
59# upload_folder(repo_id=repo_id, folder_path=save_path, repo_type='model')
