Views
No views yet
1
2from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, LlamaTokenizer, GenerationConfig, pipeline
3from peft import PeftModel, PeftMixedModel
4import torch
5
6
7model_name = "0xtaipoian/open-lilm"
8
9bnb_config = BitsAndBytesConfig(
10 load_in_4bit=True,
11 bnb_4bit_use_double_quant=True,
12 bnb_4bit_quant_type="nf4",
13 bnb_4bit_compute_dtype=torch.bfloat16
14)
15tokenizer = AutoTokenizer.from_pretrained(model_name)
16model = AutoModelForCausalLM.from_pretrained(
17 model_name,
18 torch_dtype=torch.bfloat16,
19 trust_remote_code=True,
20 quantization_config=bnb_config,
21)
22
23def chat(messages, temperature=0.9, max_new_tokens=200):
24 input_ids = tokenizer.apply_chat_template(conversation=messages, tokenize=True, add_generation_prompt=True, return_tensors='pt').to('cuda:0')
25 output_ids = model.generate(input_ids, max_new_tokens=max_new_tokens, temperature=temperature, do_sample=True)
26
27 chatml = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
28 print(chatml)
29
30 response = tokenizer.decode(output_ids[0][input_ids.shape[1]:], skip_special_tokens=False)
31
32 return response
33
34messages = [
35 # {"role": "system", "content": ""},
36 {"role": "user",
37
38 "content":
39"""
40密陽44人輪姦案」受害女隔20年現身:時間停在2004,不記得
41"""}]
42
43result = chat(messages, max_new_tokens=200, temperature=1)
44
45print(result)