Views
No views yet
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
3
4bnb_config = BitsAndBytesConfig(
5 load_in_4bit=True,
6 bnb_4bit_use_double_quant=True,
7 bnb_4bit_quant_type="nf4",
8 bnb_4bit_compute_dtype=torch.bfloat16
9)
10
11tokenizer = AutoTokenizer.from_pretrained("LuuNgoc2k2/Law-Llama-v1", add_eos_token=True, padding_side='right')
12model = AutoModelForCausalLM.from_pretrained(
13 'LuuNgoc2k2/Law-Llama-v1',
14 torch_dtype=torch.bfloat16,
15 quantization_config=bnb_config, # If you need
16 device_map="auto",
17 use_cache=True,
18)
19
20tokenizer.pad_token = tokenizer.eos_token
211PROMPT = """
2### Hướng dẫn: Bạn là một trợ lí Tiếng Việt. Hãy luôn trả lời một cách trung thực và an toàn
3Câu trả lời của bạn không nên chứa bất kỳ nội dung gây hại, nguy hiểm hoặc bất hợp pháp nào
4Nếu một câu hỏi không có ý nghĩa hoặc không hợp lý về mặt thông tin, hãy giải thích tại sao thay vì trả lời một điều gì đó không chính xác
5Nếu bạn không biết câu trả lời cho một câu hỏi, hãy trẳ lời là bạn không biết và vui lòng không chia sẻ thông tin sai lệch.
6### Câu hỏi: {input}
7"""
8
9question = """Trình bày về thủ tục li hôn ?"""
10
11text = PROMPT.format_map({
12 'input': question,
13})
14
15
16input_ids = tokenizer(text, return_tensors='pt', add_special_tokens=False).to('cuda')
17
18generated_ids = model.generate(
19 input_ids=input_ids['input_ids'],
20 max_new_tokens=1024,
21 do_sample=True,
22 top_p=0.95,
23 top_k=40,
24 temperature=0.3,
25 repetition_penalty=1.1,
26 no_repeat_ngram_size=7,
27 num_beams=5,
28)
29
30a = tokenizer.batch_decode(generated_ids)[0]
31# print(a.split('### Trả lời:')[1])
32print(a)