Views
No views yet
1from huggingface_hub import create_repo, upload_folder
2import os
3
4import torch
5import transformers
6from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig, pipeline, set_seed
7
8model_id = "nvidia/Hymba-1.5B-Instruct"
9repo_id = "yujiepan/hymba-tiny-random"
10save_path = f"/tmp/{repo_id}"
11
12config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
13config.conv_dim = {str(i): 32 for i in range(3)}
14config.hidden_size = 16
15config.intermediate_size = 32
16config.num_attention_heads = 2
17config.num_key_value_heads = 1
18config.v_head_dim = 8
19config.num_hidden_layers = 3
20
21tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
22tokenizer.save_pretrained(save_path)
23
24model = AutoModelForCausalLM.from_config(
25 config, torch_dtype=torch.bfloat16, trust_remote_code=True,
26)
27model.generation_config = GenerationConfig.from_pretrained(
28 model_id, trust_remote_code=True)
29
30set_seed(42)
31with torch.no_grad():
32 for _, p in sorted(model.named_parameters()):
33 torch.nn.init.uniform_(p, -0.2, 0.2)
34
35model.save_pretrained(save_path)
36
37prompt = 'Hello!'
38messages = [
39 {"role": "system", "content": "You are a helpful assistant."}
40]
41messages.append({"role": "user", "content": prompt})
42tokenized_chat = tokenizer.apply_chat_template(
43 messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to('cuda')
44outputs = model.cuda().generate(
45 tokenized_chat,
46 max_new_tokens=16,
47 do_sample=False,
48 temperature=0.7,
49 use_cache=True,
50)
51input_length = tokenized_chat.shape[1]
52response = tokenizer.decode(
53 outputs[0][input_length:], skip_special_tokens=True)
54print(f"Model response: {response}")
55
56os.system(f"ls -alh {save_path}")
57create_repo(repo_id, exist_ok=True)
58upload_folder(repo_id=repo_id, folder_path=save_path)