Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2import torch
3
4def generate_response(input):
5 '''
6 simple test for the model
7 '''
8 # tokenzize the input
9 tokenized_input = tokenizer.encode_plus(input, return_tensors='pt').to(device)
10
11 # generate the response
12 outputs = model.generate(
13 input_ids=tokenized_input['input_ids'],
14 attention_mask=tokenized_input['attention_mask'],
15 pad_token_id=tokenizer.pad_token_id,
16 do_sample=False,
17 repetition_penalty=1.3,
18 max_length=500
19 )
20
21 # decode the response
22 return tokenizer.decode(outputs[0], skip_special_tokens=True)
23
24if __name__ == '__main__':
25 device = 'cuda' if torch.cuda.is_available() else 'cpu'
26 model = AutoModelForCausalLM.from_pretrained("DavidLanz/Taiwan-tinyllama-v1.0-chat",device_map=device,torch_dtype=torch.bfloat16)
27 tokenizer = AutoTokenizer.from_pretrained("DavidLanz/Taiwan-tinyllama-v1.0-chat")
28 while(True):
29 text = input("input a simple prompt:")
30 print('System:', generate_response(text))