Views
No views yet
1from huggingface_hub import create_repo, upload_folder
2from transformers import (
3 pipeline,
4 set_seed,
5 AutoConfig,
6 AutoModelForCausalLM,
7 AutoTokenizer,
8 GenerationConfig,
9)
10import torch
11import transformers
12import os
13
14model_id = "HuggingFaceTB/SmolLM-1.7B"
15repo_id = "yujiepan/smollm-tiny-random"
16save_path = f"/tmp/{repo_id}"
17
18config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
19config._name_or_path = model_id
20config.hidden_size = 8
21config.intermediate_size = 16
22config.num_attention_heads = 4
23config.num_hidden_layers = 2
24config.num_key_value_heads = 2
25print(config)
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, attn_implementation="eager", trust_remote_code=True
32)
33model.generation_config = GenerationConfig.from_pretrained(model_id, trust_remote_code=True)
34
35set_seed(42)
36with torch.no_grad():
37 for _, p in sorted(model.named_parameters()):
38 torch.nn.init.uniform_(p, -0.1, 0.1)
39
40model.save_pretrained(save_path)
41
42pipe = pipeline("text-generation", model=save_path, device="cuda", trust_remote_code=True)
43print(pipe("Hello World!"))
44
45# messages = [
46# {"role": "system", "content": "You are a robot."},
47# {"role": "user", "content": "Hi!"},
48# ]
49# chatbot = pipeline("text-generation", model=save_path, max_length=1000, max_new_tokens=16, trust_remote_code=True)
50# print(chatbot(messages))