Views
No views yet
1from transformers import pipeline
2from huggingface_hub import create_repo, upload_folder
3import torch
4import transformers
5import os
6
7model_id = 'mistralai/Mistral-Nemo-Instruct-2407'
8repo_id = 'yujiepan/mistral-nemo-2407-tiny-random'
9save_path = f'/tmp/{repo_id}'
10
11config = transformers.AutoConfig.from_pretrained(model_id)
12config.hidden_size = 8
13config.intermediate_size = 32
14config.num_attention_heads = 4
15config.num_hidden_layers = 2
16config.num_key_value_heads = 2
17config.head_dim = 2
18print(config)
19
20tokenizer = transformers.AutoTokenizer.from_pretrained(model_id)
21tokenizer.save_pretrained(save_path)
22
23model = transformers.AutoModelForCausalLM.from_config(config, torch_dtype=torch.bfloat16)
24model.generation_config = transformers.GenerationConfig.from_pretrained(model_id)
25
26transformers.set_seed(42)
27with torch.no_grad():
28 for _, p in sorted(model.named_parameters()):
29 torch.nn.init.uniform_(p, -0.1, 0.1)
30
31pipe = pipeline('text-generation', model=model, tokenizer=tokenizer, do_sample=False, device='cuda')
32print(pipe('Hello World!'))
33
34messages = [
35 {"role": "system", "content": "You are a pirate chatbot who always responds in pirate speak!"},
36 {"role": "user", "content": "Who are you?"},
37]
38chatbot = pipeline("text-generation", model=save_path, max_length=1000, max_new_tokens=16)
39print(chatbot(messages))
40
41model.save_pretrained(save_path)