Views
No views yet

1from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
2import torch
3from transformers import TextStreamer
4
5def generate_response(input):
6 '''
7 simple test for the model
8 '''
9 # tokenzize the input
10 tokenized_input = tokenizer.encode_plus(input, return_tensors='pt').to(device)
11 print(tokenized_input['input_ids'])
12 # generate the response
13 _ = model.generate(
14 input_ids=tokenized_input['input_ids'],
15 attention_mask=tokenized_input['attention_mask'],
16 pad_token_id=tokenizer.pad_token_id,
17 do_sample=True,
18 repetition_penalty=1.0,
19 max_length=2048,
20 streamer=streamer,
21 )
22
23
24if __name__ == '__main__':
25 device = 'cuda' if torch.cuda.is_available() else 'cpu'
26 model = AutoModelForCausalLM.from_pretrained("benchang1110/Taiwan-tinyllama-v1.1-base",attn_implementation="flash_attention_2",device_map=device,torch_dtype=torch.bfloat16)
27 tokenizer = AutoTokenizer.from_pretrained("benchang1110/Taiwan-tinyllama-v1.1-base",use_fast=True)
28 streamer = TextStreamer(tokenizer)
29 while(True):
30 text = input("input a simple prompt:")
31 generate_response(text)| Data size | Global Batch Size | Learning Rate | Epochs | Max Length | Weight Decay |
|---|---|---|---|---|---|
| 10B | 32 | 5e-5 | 1 | 2048 | 1e-4 |
