Hugging Face's logo

ngxson
/
Vistral-7B-ChatML

Text Generation
Transformers
Safetensors
GGUF
Vietnamese
mistral
LLMs
NLP
Vietnamese
conversational
Inference Endpoints
text-generation-inference
Model card Files Community
Vistral-7B-ChatML
File size: 4,527 Bytes 
ae1a83e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
 
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, HfArgumentParser, TrainingArguments, pipeline, logging, TextStreamer
from peft import LoraConfig, PeftModel, prepare_model_for_kbit_training, get_peft_model
import os, torch, wandb, platform, warnings
from datasets import load_dataset
from trl import SFTTrainer

hf_token = ''
wnb_token = ''
wnb_name = 'vistral-chatml'
MODEL = 'Viet-Mistral/Vistral-7B-Chat'
resume_from_checkpoint = False
output_dir = 'vistral-chatml'
tokenizer_path = '.'

#######################################################
## DATASET


from datasets import load_dataset


def generate_system_prompt(i):
    system_prompt = "Bạn là một trợ lí Tiếng Việt nhiệt tình và trung thực. Hãy luôn trả lời một cách hữu ích nhất có thể, đồng thời giữ an toàn."
    if i % 2 == 0:
        system_prompt += "\nCâu trả lời của bạn không nên chứa bất kỳ nội dung gây hại, phân biệt chủng tộc, phân biệt giới tính, độc hại, nguy hiểm hoặc bất hợp pháp nào. Hãy đảm bảo rằng các câu trả lời của bạn không có thiên kiến xã hội và mang tính tích cực."
    if i % 5 == 0:
        system_prompt += "\nNế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. Nế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."
    return system_prompt

tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)

def tokenize_chat(input, i):
    print(generate_system_prompt(i))
    conversation = [{'role': 'system', 'content': generate_system_prompt(i)}]
    for msg in input['conversations']:
        output = {'role': 'user', 'content': msg['value']}
        if msg['from'] == 'gpt':
            output['role'] = 'assistant'
        conversation.append(output)
    formatted = tokenizer.apply_chat_template(conversation, tokenize=False)
    return tokenizer(formatted)

sharegpt_dataset = load_dataset('bkai-foundation-models/vi-self-chat-sharegpt-format')
train_data = sharegpt_dataset['train'].shuffle(seed=42)\
    .select(range(800))\
    .map(lambda x, i: tokenize_chat(x, i), remove_columns=["conversations"], with_indices=True)


#######################################################
## SETUP

wandb.login(key=wnb_token)
wandb.init(name=wnb_name)
# use custom tokenizer instead of one comes from the model
#tokenizer = AutoTokenizer.from_pretrained(
#  MODEL,
#  add_eos_token=False,
#  add_bos_token=False,
#  token=hf_token,
#)
bnb_config = BitsAndBytesConfig(
  load_in_4bit=True,
  bnb_4bit_quant_type="nf4",
  bnb_4bit_compute_dtype=torch.bfloat16,
  bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(
  MODEL,
  device_map="auto",
  token=hf_token,
  quantization_config=bnb_config,
  trust_remote_code=True,
)


#######################################################
## LORA CONFIG

model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)
peft_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
        "lm_head",
    ],
    bias="none",
    lora_dropout=0.05,  # Conventional
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()

from accelerate import Accelerator
accelerator = Accelerator()
model = accelerator.prepare_model(model)


#######################################################
## TRAIN

from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling
trainer = Trainer(
    model=model,
    train_dataset=train_data,
    args=TrainingArguments(
        report_to='wandb',
        warmup_steps=1,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        gradient_checkpointing=True,
        num_train_epochs=4,
        learning_rate=2.5e-5,
        logging_steps=1,
        optim="paged_adamw_8bit",
        save_strategy="steps",
        save_steps=10,
        save_total_limit=4,
        output_dir=output_dir
    ),
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
model.config.use_cache = False

trainer.train(resume_from_checkpoint=resume_from_checkpoint)