machine-translation / llm_toolkit /translation_engine.py
dh-mc's picture
initial code for Chinese/English translation
3860729
raw
history blame
3.79 kB
import os
import pandas as pd
import torch
from unsloth import FastLanguageModel, is_bfloat16_supported
from trl import SFTTrainer
from transformers import TrainingArguments, TextStreamer
from llm_toolkit.translation_utils import *
from llamafactory.chat import ChatModel
print(f"loading {__file__}")
def get_model_names(
model_name, save_method="merged_4bit_forced", quantization_method="q5_k_m"
):
hub_model = model_name.split("/")[-1] + "-MAC-"
local_model = "models/" + hub_model
return {
"local": local_model + save_method,
"local-gguf": local_model + quantization_method,
"hub": hub_model + save_method,
"hub-gguf": hub_model + "gguf-" + quantization_method,
}
def load_model(
model_name,
max_seq_length=2048,
dtype=None,
load_in_4bit=False,
template="chatml",
adapter_name_or_path=None,
):
print(f"loading model: {model_name}")
if adapter_name_or_path:
args = dict(
model_name_or_path=model_name,
adapter_name_or_path=adapter_name_or_path, # load the saved LoRA adapters
template=template, # same to the one in training
finetuning_type="lora", # same to the one in training
quantization_bit=4, # load 4-bit quantized model
)
chat_model = ChatModel(args)
return chat_model.engine.model, chat_model.engine.tokenizer
model, tokenizer = FastLanguageModel.from_pretrained(
model_name=model_name, # YOUR MODEL YOU USED FOR TRAINING
max_seq_length=max_seq_length,
dtype=dtype,
load_in_4bit=load_in_4bit,
trust_remote_code=True,
)
FastLanguageModel.for_inference(model)
return model, tokenizer
def test_model(model, tokenizer, prompt):
inputs = tokenizer(
[prompt],
return_tensors="pt",
).to("cuda")
text_streamer = TextStreamer(tokenizer)
_ = model.generate(
**inputs, max_new_tokens=128, streamer=text_streamer, use_cache=True
)
def load_trainer(
model,
tokenizer,
dataset,
num_train_epochs,
max_seq_length=2048,
fp16=False,
bf16=False,
output_dir="./outputs",
):
model = FastLanguageModel.get_peft_model(
model,
r=16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
target_modules=[
"q_proj",
"k_proj",
"v_proj",
"o_proj",
"gate_proj",
"up_proj",
"down_proj",
],
lora_alpha=16,
lora_dropout=0, # Supports any, but = 0 is optimized
bias="none", # Supports any, but = "none" is optimized
# [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
use_gradient_checkpointing="unsloth", # True or "unsloth" for very long context
random_state=3407,
use_rslora=False, # We support rank stabilized LoRA
loftq_config=None, # And LoftQ
)
trainer = SFTTrainer(
model=model,
tokenizer=tokenizer,
train_dataset=dataset,
dataset_text_field="text",
max_seq_length=max_seq_length,
dataset_num_proc=2,
packing=False, # Can make training 5x faster for short sequences.
args=TrainingArguments(
per_device_train_batch_size=2,
gradient_accumulation_steps=4,
warmup_steps=5,
num_train_epochs=num_train_epochs,
learning_rate=2e-4,
fp16=not is_bfloat16_supported(),
bf16=is_bfloat16_supported(),
logging_steps=100,
optim="adamw_8bit",
weight_decay=0.01,
lr_scheduler_type="linear",
seed=3407,
output_dir=output_dir,
),
)
return trainer