Views
No views yet
1import transformers
2import torch
3import os
4from huggingface_hub import create_repo, upload_folder
5
6source_model_id = 'THUDM/chatglm3-6b-128k'
7tiny_random_name = 'chatglm3-tiny-random'
8save_path = f'/tmp/yujiepan/{tiny_random_name}'
9repo_id = f'yujiepan/{tiny_random_name}'
10
11config = transformers.AutoConfig.from_pretrained(
12 source_model_id, trust_remote_code=True)
13config.hidden_size = 4
14config.ffn_hidden_size = 6
15config.num_attention_heads = 4
16config.kv_channels = 2
17config.num_layers = 2
18config.torch_dtype = torch.float16
19
20model = transformers.AutoModelForCausalLM.from_config(
21 config, trust_remote_code=True, torch_dtype=torch.float16)
22model = model.half()
23
24tokenizer = transformers.AutoTokenizer.from_pretrained(
25 source_model_id, trust_remote_code=True)
26
27# result = transformers.pipelines.pipeline(
28# 'text-generation',
29# model=model, tokenizer=tokenizer,
30# device=0,
31# max_new_tokens=16,
32# )('Hello')
33# print(result)
34model = model.cuda()
35response, history = model.chat(tokenizer, "Hi", history=[], max_length=32)
36print(response)
37
38model.save_pretrained(save_path)
39tokenizer.save_pretrained(save_path)
40
41os.system(f'ls -alh {save_path}')
42create_repo(repo_id, exist_ok=True)
43upload_folder(repo_id=repo_id, folder_path=save_path)