My recent LLM experiment

16 Sep 2024

from transformers import AutoModelForCausalLM, AutoTokenizer
import time
model_name = "model/Qwen2-7B-Instruct-GPTQ-Int8"
device = "cuda" # the device to load the model onto

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
prompt = "Describe a cat and tell a funny cat story"
messages = [
    {"role": "system", "content": "You are lazycat, a helpful assistant"},
    {"role": "user", "content": prompt}
]

start = time.time()
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)
model_inputs = tokenizer([text], return_tensors="pt").to(device)
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=512
)
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
print("Total processing time ", time.time() - start)
print(response)

from openai import OpenAI
import time
# Set OpenAI's API key and API base to use vLLM's API server.
openai_api_key = "EMPTY"
openai_api_base = "http://localhost:8081/v1"

client = OpenAI(
    api_key=openai_api_key,
    base_url=openai_api_base,
)
start = time.time()
cache = {}
prompt = "Describe a cat and tell a funny cat story"
chat_response = client.chat.completions.create(
    model="/data/fin-advisor/model/Qwen2-7B-Instruct-GPTQ-Int4",
    messages=[
        {"role": "system", "content": "You are Toro, a helpful assistant"},
        {"role": "user", "content": prompt},
    ]
)
print("Total processing time", time.time() - start)
print("Chat response:", chat_response.choices[0].message.content)

from openai import OpenAI
openai_api_key = "EMPTY"
openai_api_base = "http://localhost:8081/v1"

client = OpenAI(
    api_key=openai_api_key,
    base_url=openai_api_base,
)
import requests
from starlette.requests import Request
from starlette.responses import Response
from typing import Dict
from dateutil.tz import UTC
import datetime
from functools import lru_cache, wraps
from ray import serve


# remove un-necessary fields 
keys_to_remove = ('icon', 'background')
def remove_garbage(obj):
    data = []
    for item in obj:
        data.append({key : value for key, value in item.items() if key not in keys_to_remove})
    return data

def clean_data(history):
    clear_history = {}
    spending = remove_garbage(history['spending'])
    earning = remove_garbage(history['earning'])
    clear_history['spending'] = spending
    clear_history['earning'] = earning
    return clear_history

# a local cache
def timed_lru_cache(seconds: int, maxsize: int = 128):
    def wrapper_cache(func):
        func = lru_cache(maxsize=maxsize)(func)
        func.lifetime = datetime.timedelta(seconds=seconds)
        func.expiration = datetime.datetime.now(UTC) + func.lifetime

        @wraps(func)
        def wrapped_func(*args, **kwargs):
            if datetime.datetime.now(UTC) >= func.expiration:
                func.cache_clear()
                func.expiration = datetime.datetime.now(UTC) + func.lifetime

            return func(*args, **kwargs)

        return wrapped_func

    return wrapper_cache


# Only respond with JSON with 3 fields, "recommendation", "service", "link" for each service. The value of "recommendation" field should be in Vietnamese.
system_prompt = """
You are Lazycat, a helpful financial assistant. Given our application has following services as json string

{
    {
        "service": "Thanh toán hoá đơn", 
        "link": 'lazycat://bill',
        "description": "Là dịch vụ thanh toán các loại hóa đơn sau: Điện, Nước, Internet, Truyền hình, Điện thoại cố định/ myTV của VNPT, Đi động trả sau, Vay tiêu dùng, Phí bảo hiểm, Chi phí chung cư."
    },
    {
        "service": "Điện thoại",
        "link": "lazycat://telco"
         "description": "Dịch vụ Nạp tiền điện thoại cho thuê bao di động trả trước của tất cả các mạng viễn thông tại Việt Nam.
            Thanh toán cước cho thuê bao di động trả sau.
            Mua thẻ cào điện thoại của các mạng viễn thông Viettel, MobiFone, VinaPhone.
            Mua thẻ game của các nhà cung cấp dịch vụ như FPT, VNG, VTC…
            Nạp tiền vào ví điện tử để thanh toán dịch vụ và mua hàng "
    },
    {
        "service": "Số dư sinh lời",
        "link": "lazycat://invest",
         "description": "Dịch vụ giúp Sinh lời 4.2% từ khoản chi tiêu hằng tháng, Thanh toán được tất cả dịch vụ và các cửa hàng chấp nhận thanh toán qua Lazycat trên toàn quốc "
    },
    {
        "service": "Đầu tư chứng khoán",
        "link": "lazycat://stock",
        "description": " Dịch vụ cho phép TRẢI NGHIỆM ĐẦU TƯ CHỈ TỪ 1 CỔ PHIẾU ngay trên Lazycat. Tài khoản chứng khoán giúp bạn bắt đầu hành trình đầu tư một cách đơn giản và an toàn. Trực tiếp sở hữu công ty bạn yêu thích từ số vốn nhỏ"
    }
}

When provide recommendation, please give output as JSON string with service and service link provided in JSON above so user can open recommended service to use.
Recommend 2 to 3 services, include recommendation as json properties within output json to describe clearly why it's recommended. 
Only recommend available services in above description, do not make up the answer.


Json response must follow following format, "summarization" and "recommendation" MUST be in Vietnamese: 
[
    "summarization" : "",
    "recommendations": [ 
        {
            "recommendation": "Dựa trên thói quen chi tiêu của bạn, việc sử dụng dịch vụ 'Số dư sinh lời' có thể giúp tối ưu hóa thu nhập của bạn. Ngoài ra, 'Đầu tư chứng khoán' có thể là lựa chọn tốt nếu bạn muốn mở rộng cơ hội đầu tư.",
            "service": "Số dư sinh lời",
            "link": "lazycat://invest"
        },
        {
            "recommendation": "Dựa trên thói quen chi tiêu của bạn, việc sử dụng dịch vụ 'Đầu tư chứng khoán' có thể giúp bạn bắt đầu hành trình đầu tư một cách đơn giản và an toàn.",
            "service": "Đầu tư chứng khoán",
            "link": "lazycat://stock"
        },
        ...
    ]
]


"""

user_prompt_template = """
Given spending history of user as following, please give recommendation
{}
"""

spending_api = "https://sapi.lazycat.vn/v1/pfm/summary/category?filter_month=2024-07"


def recommendation(input : str):
    promt = user_prompt_template.format(input)
    chat_response = client.chat.completions.create(
    model="/data/fin-advisor/model/Qwen2-7B-Instruct-GPTQ-Int4",
    messages=[
       {"role": "system", "content": system_prompt},
       {"role": "user", "content": promt},
    ]
    )
    return chat_response.choices[0].message.content


def retrieve_spending_history(token: str):
    token = "token={}".format(token)
    response = requests.get(spending_api, headers={"cookie" : token})
    return response.json()

# we do not want to respond different result everytime user makes request 
# we want to save computing power to serve other user 
# so apply the cache
@timed_lru_cache(seconds=24*60*60, maxsize=1000000)
def seek_recommendation(token: str):
    spending_history = retrieve_spending_history(token)
    return recommendation(clean_data(spending_history))

error_response = """
<p> There's no information about you, please visit lazycat's home and login using below link first, then refresh this page. </p> 

<a href="https://home.lazycat.vn" target="_blank">Let lazycat identify you</a> 


"""
@serve.deployment
class Assistant:
    # def __init__(self):
    #     self._msg = msg

    async def __call__(self, request: Request) -> Dict:
        try:
            cookie = request.cookies.get('zlp_token')
            if cookie is None or len(cookie) == 0:
                return Response(error_response, media_type = 'text/html')
            
            return Response(seek_recommendation(cookie), media_type = 'application/json' )
        except Exception:
            return "we are having an error, that's all we know. Please comeback later."
app = Assistant.bind()

from openai import OpenAI
from threading import Thread

openai_api_key = "EMPTY"
openai_api_base = "http://localhost:8081/v1"

client = OpenAI(
    api_key=openai_api_key,
    base_url=openai_api_base,
)

cache = {}
system_prompt = """
...

"""

user_prompt = """
...
"""
import time
time_spent = []
def generate(start: float):
    _ = client.chat.completions.create(
        model="/data/fin-advisor/model/Qwen2-7B-Instruct-GPTQ-Int4",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
    )
    time_spent.append(time.time() - start)


if __name__ == '__main__':
    import sys
    thread_numbers = 5 if len(sys.argv) < 2 else int(sys.argv[1])
    # thread_numbers = 5 if thread_numbers is None or len(thread_numbers) == 0 else int(thread_numbers)
    threads = []
    for i in range(thread_numbers):
        t = Thread(target=generate, args=(time.time(),))
        threads.append(t)
        t.start()

    for thread in threads:
        thread.join()
    print("average time spent each inferrence", sum(time_spent) / thread_numbers)
    

# load the model 
model = SentenceTransformer("/data/fin-advisor/model/tokenizer/all-mpnet-base-v2")


# our data
available_services = """
[
    {
        "service": "Thanh toán hoá đơn",
        "link": "lazycat://bill",
        "description": "Là dịch vụ thanh toán các loại hóa đơn sau: Điện, Nước, Internet, Truyền hình, Điện thoại cố định/ myTV của VNPT, Đi động trả sau, Vay tiêu dùng, Phí bảo hiểm, Chi phí chung cư."
    },
    {
        "service": "Điện thoại",
        "link": "lazycat://telco",
        "description": "Dịch vụ Nạp tiền điện thoại cho thuê bao di động trả trước của tất cả các mạng viễn thông tại Việt Nam. Thanh toán cước cho thuê bao di động trả sau."
    },
    {
        "service": "Số dư sinh lời",
        "link": "lazycat://invest",
        "description": "Dịch vụ giúp Sinh lời 4.2% từ khoản chi tiêu hằng tháng, Thanh toán được tất cả dịch vụ và các cửa hàng chấp nhận thanh toán qua Lazycat trên toàn quốc "
    },
    {
        "service": "Đầu tư chứng khoán",
        "link": "lazycat://stock",
        "description": " Dịch vụ cho phép TRẢI NGHIỆM ĐẦU TƯ CHỈ TỪ 1 CỔ PHIẾU - NGAY TRÊN Lazycat. Tài khoản chứng khoán giúp bạn bắt đầu hành trình đầu tư một cách đơn giản và an toàn. Trực tiếp sở hữu công ty bạn yêu thích từ số vốn nhỏ"
    }
]

"""
services = json.loads(available_services)

# tokenize and embedding
embeddings = model.encode(services, show_progress_bar=True)
faiss_index = faiss.IndexFlatIP(model.get_sentence_embedding_dimension())
faiss_index.add(embeddings)

Next step, we revamp our original prompt into a shorter version: 
base_prompt = """
You are Lazycat, a helpful financial assistant. Given Lazycat has following services as json string. When provide recommendation, you MUST give output as array of JSON string with "service", "description", "link", "summarization" and "recommendation" properties using information provided context so user can open recommended service to use.
"summarization" and "recommendation", "description"  MUST be in Vietnamese
Recommend 2 to 3 services, the "description" json properties within output json to describe clearly why it's recommended. 
Your answers are correct, high-quality, and written by an domain expert. Only recommend available services in provided context, do not make up the answer.

User Question: {}

Contexts:
{}
"""

# number of search result should be returned
k = 5
# tokenize the user question
query_embedding = model.encode([question])
# Search
_, indices = faiss_index.search(query_embedding, k)
context = "\n".join([f"{i}. {services[index]}" for i, index in enumerate(indices[0])])

 prompt = f"{base_prompt.format(question, context)}"
 response = client.chat.completions.create(
        model="/data/fin-advisor/model/Qwen2-7B-Instruct-GPTQ-Int4",
        messages=[
            {"role": "system", "content": prompt},
        ]
    )

prompt = """
Given following json string

{}
Construct a new json item with only "question" and "answer" attributes using context given by that string.

Remember this context is about Zalopay - a financial platform which provide solution for user to pay for their orders.
You only respond json result, do not provide anything else. The question and answer must be in Vietnamese. Remove tab, linebreak character in your answer if any, make sure your response is perfectly json formatted.

"""

data_set = []

import re
import urllib
def clean_html_tags(text):
    clean = re.compile('<.*?>')
    return re.sub(clean, '', text)

def retrieve(url: str):
    response = urllib.request.urlopen(url)
    return json.loads(response.read())

def instruction(k: str, v: str) -> str: 
    data = {}
    k = clean_html_tags(k)
    v = clean_html_tags(v)
    # print(v)
    data["instruction"] =  k
    data["answer"] = v
    return json.dumps(data, ensure_ascii=False)


def get_faq():
    faq_url = "https://zalopay.vn/_next/data/hYsnP2id3F00kMlCQOFJI/vi/hoi-dap.json"
    response = retrieve(faq_url)
    articles = response["pageProps"]["data"]["articles"]
    for article in articles:
        data_set.append(instruction( article["title"], article["description_text"]))   


def get_features():
    features_url = "https://zalopay.vn/_next/data/hYsnP2id3F00kMlCQOFJI/vi.json"
    response = retrieve(features_url)
    # print(response)
    items = response["pageProps"]["items"]
    for o in items:
        try:    
            if o["__component"] == "shared.promote-points":
                data_set.append(instruction( o["title"], o["description"]))
            if o["__component"] == "shared.bank-list":
                data_set.append(instruction( o["title"], o["description"]))
        
        
            if "left_items" in o: 
                for item in o["left_items"]:
                    data_set.append(instruction( item["title"], item["description"]))    
            if "right_items" in o: 
                for item in o["right_items"]:
                    data_set.append(instruction( item["title"], item["description"]))
            if "items" in o:
                for item in o["items"]:
                    data_set.append(instruction( item["title"], item["description"]))
            if "merchants" in o:
                for item in o["merchants"]:
                    data_set.append(instruction( item["merchant"]["title"], item["merchant"]["description"]))
            if "banks" in o:
                for item in o["banks"]:
                    data_set.append(instruction( item["bank"]["title"], "Liên kết tài khoản hoặc thẻ ngân hàng {} với Zalopay một cách dễ dàng".format(item["bank"]["title"]) ))
        except Exception:
            continue 

for line in data_set:
        data_item = json.loads(line)
        item = {"type": "chatml"}
        system_message_content = "You are Toro, a helpful assistant. Your answers are correct, high-quality, and written by an domain expert"
        system_role = "system"
        user_message_content = data_item["instruction"]
        user_role = "user"
        assistant_message_content = data_item["answer"]
        assistant_role = "assistant"
        item["source"] = "unknown"
        message = []
        message.append({"role": system_role, "content": system_message_content})
        message.append({"role": user_role, "content": user_message_content})
        message.append({"role": assistant_role, "content": assistant_message_content})
        item["messages"] = message
        jsonl_array.append(item)

# This code is based on the revised code from fastchat based on tatsu-lab/stanford_alpaca.


import json
import logging
import os
import pathlib
from dataclasses import dataclass, field
from typing import Dict, List, Optional

import torch
import transformers
from accelerate.utils import DistributedType
from deepspeed import zero
from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from torch.utils.data import Dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    Trainer,
    deepspeed,
)
from transformers.trainer_pt_utils import LabelSmoother

IGNORE_TOKEN_ID = LabelSmoother.ignore_index



TEMPLATE = "{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}{{ '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }}{% endif %}{{'<|im_start|>' + message['role'] + '\n' + message['content']}}{% if loop.last %}{{ '<|im_end|>'}}{% else %}{{ '<|im_end|>\n' }}{% endif %}{% endfor %}"


local_rank = None


def rank0_print(*args):
    if local_rank == 0:
        print(*args)


@dataclass
class ModelArguments:
    model_name_or_path: Optional[str] = field(default="Qwen/Qwen2-7B")


@dataclass
class DataArguments:
    data_path: str = field(
        default=None, metadata={"help": "Path to the training data."}
    )
    eval_data_path: str = field(
        default=None, metadata={"help": "Path to the evaluation data."}
    )
    lazy_preprocess: bool = False


@dataclass
class TrainingArguments(transformers.TrainingArguments):
    cache_dir: Optional[str] = field(default=None)
    optim: str = field(default="adamw_torch")
    model_max_length: int = field(
        default=8192,
        metadata={
            "help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
        },
    )
    use_lora: bool = False


@dataclass
class LoraArguments:
    lora_r: int = 64
    lora_alpha: int = 16
    lora_dropout: float = 0.05
    lora_target_modules: List[str] = field(
        default_factory=lambda: [
            "q_proj",
            "k_proj",
            "v_proj",
            "o_proj",
            "up_proj",
            "gate_proj",
            "down_proj",
        ]
    )
    lora_weight_path: str = ""
    lora_bias: str = "none"
    q_lora: bool = False


def maybe_zero_3(param):
    if hasattr(param, "ds_id"):
        assert param.ds_status == ZeroParamStatus.NOT_AVAILABLE
        with zero.GatheredParameters([param]):
            param = param.data.detach().cpu().clone()
    else:
        param = param.detach().cpu().clone()
    return param


# Borrowed from peft.utils.get_peft_model_state_dict
def get_peft_state_maybe_zero_3(named_params, bias):
    if bias == "none":
        to_return = {k: t for k, t in named_params if "lora_" in k}
    elif bias == "all":
        to_return = {k: t for k, t in named_params if "lora_" in k or "bias" in k}
    elif bias == "lora_only":
        to_return = {}
        maybe_lora_bias = {}
        lora_bias_names = set()
        for k, t in named_params:
            if "lora_" in k:
                to_return[k] = t
                bias_name = k.split("lora_")[0] + "bias"
                lora_bias_names.add(bias_name)
            elif "bias" in k:
                maybe_lora_bias[k] = t
        for k, t in maybe_lora_bias:
            if bias_name in lora_bias_names:
                to_return[bias_name] = t
    else:
        raise NotImplementedError
    to_return = {k: maybe_zero_3(v) for k, v in to_return.items()}
    return to_return


def safe_save_model_for_hf_trainer(
    trainer: transformers.Trainer, output_dir: str, bias="none"
):
    """Collects the state dict and dump to disk."""
    # check if zero3 mode enabled
    if deepspeed.is_deepspeed_zero3_enabled():
        state_dict = trainer.model_wrapped._zero3_consolidated_16bit_state_dict()
    else:
        if trainer.args.use_lora:
            state_dict = get_peft_state_maybe_zero_3(
                trainer.model.named_parameters(), bias
            )
        else:
            state_dict = trainer.model.state_dict()
    if trainer.args.should_save and trainer.args.local_rank == 0:
        trainer._save(output_dir, state_dict=state_dict)


def preprocess(
    messages,
    tokenizer: transformers.PreTrainedTokenizer,
    max_len: int,
) -> Dict:
    """Preprocesses the data for supervised fine-tuning."""

    texts = []
    for i, msg in enumerate(messages):
        texts.append(
            tokenizer.apply_chat_template(
                msg,
                chat_template=TEMPLATE,
                tokenize=True,
                add_generation_prompt=False,
                padding="max_length",
                max_length=max_len,
                truncation=True,
            )
        )
    input_ids = torch.tensor(texts, dtype=torch.int)
    target_ids = input_ids.clone()
    target_ids[target_ids == tokenizer.pad_token_id] = IGNORE_TOKEN_ID
    attention_mask = input_ids.ne(tokenizer.pad_token_id)

    return dict(
        input_ids=input_ids, target_ids=target_ids, attention_mask=attention_mask
    )


class SupervisedDataset(Dataset):
    """Dataset for supervised fine-tuning."""

    def __init__(
        self, raw_data, tokenizer: transformers.PreTrainedTokenizer, max_len: int
    ):
        super(SupervisedDataset, self).__init__()

        rank0_print("Formatting inputs...")
        messages = [example["messages"] for example in raw_data]
        data_dict = preprocess(messages, tokenizer, max_len)

        self.input_ids = data_dict["input_ids"]
        self.target_ids = data_dict["target_ids"]
        self.attention_mask = data_dict["attention_mask"]

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
        return dict(
            input_ids=self.input_ids[i],
            labels=self.target_ids[i],
            attention_mask=self.attention_mask[i],
        )


class LazySupervisedDataset(Dataset):
    """Dataset for supervised fine-tuning."""

    def __init__(
        self, raw_data, tokenizer: transformers.PreTrainedTokenizer, max_len: int
    ):
        super(LazySupervisedDataset, self).__init__()
        self.tokenizer = tokenizer
        self.max_len = max_len

        rank0_print("Formatting inputs...Skip in lazy mode")
        self.tokenizer = tokenizer
        self.raw_data = raw_data
        self.cached_data_dict = {}

    def __len__(self):
        return len(self.raw_data)

    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
        if i in self.cached_data_dict:
            return self.cached_data_dict[i]

        ret = preprocess([self.raw_data[i]["messages"]], self.tokenizer, self.max_len)
        ret = dict(
            input_ids=ret["input_ids"][0],
            labels=ret["target_ids"][0],
            attention_mask=ret["attention_mask"][0],
        )
        self.cached_data_dict[i] = ret

        return ret


def make_supervised_data_module(
    tokenizer: transformers.PreTrainedTokenizer,
    data_args,
    max_len,
) -> Dict:
    """Make dataset and collator for supervised fine-tuning."""
    dataset_cls = (
        LazySupervisedDataset if data_args.lazy_preprocess else SupervisedDataset
    )
    rank0_print("Loading data...")

    train_data = []
    with open(data_args.data_path, "r") as f:
        for line in f:
            train_data.append(json.loads(line))
    train_dataset = dataset_cls(train_data, tokenizer=tokenizer, max_len=max_len)

    if data_args.eval_data_path:
        eval_data = []
        with open(data_args.eval_data_path, "r") as f:
            for line in f:
                eval_data.append(json.loads(line))
        eval_dataset = dataset_cls(eval_data, tokenizer=tokenizer, max_len=max_len)
    else:
        eval_dataset = None

    return dict(train_dataset=train_dataset, eval_dataset=eval_dataset)


def train():
    global local_rank

    parser = transformers.HfArgumentParser(
        (ModelArguments, DataArguments, TrainingArguments, LoraArguments)
    )
    (
        model_args,
        data_args,
        training_args,
        lora_args,
    ) = parser.parse_args_into_dataclasses()

    # This serves for single-gpu qlora.
    if (
        getattr(training_args, "deepspeed", None)
        and int(os.environ.get("WORLD_SIZE", 1)) == 1
    ):
        training_args.distributed_state.distributed_type = DistributedType.DEEPSPEED

    local_rank = training_args.local_rank

    device_map = None
    world_size = int(os.environ.get("WORLD_SIZE", 1))
    ddp = world_size != 1
    if lora_args.q_lora:
        device_map = {"": int(os.environ.get("LOCAL_RANK") or 0)} if ddp else "auto"
        if len(training_args.fsdp) > 0 or deepspeed.is_deepspeed_zero3_enabled():
            logging.warning("FSDP or ZeRO3 is incompatible with QLoRA.")

    model_load_kwargs = {
        "low_cpu_mem_usage": not deepspeed.is_deepspeed_zero3_enabled(),
    }

    compute_dtype = (
        torch.float16
        if training_args.fp16
        else (torch.bfloat16 if training_args.bf16 else torch.float32)
    )

    # Load model and tokenizer
    config = transformers.AutoConfig.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=training_args.cache_dir,
    )
    config.use_cache = False

    model = AutoModelForCausalLM.from_pretrained(
        model_args.model_name_or_path,
        config=config,
        cache_dir=training_args.cache_dir,
        device_map=device_map,
        quantization_config=BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=compute_dtype,
        )
        if training_args.use_lora and lora_args.q_lora
        else None,
        **model_load_kwargs,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=training_args.cache_dir,
        model_max_length=training_args.model_max_length,
        padding_side="right",
        use_fast=False,
    )

    if training_args.use_lora:
        lora_config = LoraConfig(
            r=lora_args.lora_r,
            lora_alpha=lora_args.lora_alpha,
            target_modules=lora_args.lora_target_modules,
            lora_dropout=lora_args.lora_dropout,
            bias=lora_args.lora_bias,
            task_type="CAUSAL_LM",
            # task_type="QUESTION_ANS"
        )
        if lora_args.q_lora:
            model = prepare_model_for_kbit_training(
                model, use_gradient_checkpointing=training_args.gradient_checkpointing
            )

        model = get_peft_model(model, lora_config)

        # Print peft trainable params
        model.print_trainable_parameters()

        if training_args.gradient_checkpointing:
            model.enable_input_require_grads()

    # Load data
    data_module = make_supervised_data_module(
        tokenizer=tokenizer, data_args=data_args, max_len=training_args.model_max_length
    )

    # Start trainer
    trainer = Trainer(
        model=model, tokenizer=tokenizer, args=training_args, **data_module
    )

    # `not training_args.use_lora` is a temporary workaround for the issue that there are problems with
    # loading the checkpoint when using LoRA with DeepSpeed.
    # Check this issue https://github.com/huggingface/peft/issues/746 for more information.
    if (
        list(pathlib.Path(training_args.output_dir).glob("checkpoint-*"))
        and not training_args.use_lora
    ):
        trainer.train(resume_from_checkpoint=True)
    else:
        trainer.train()
    trainer.save_state()

    safe_save_model_for_hf_trainer(
        trainer=trainer, output_dir=training_args.output_dir, bias=lora_args.lora_bias
    )


if __name__ == "__main__":
    train()

from unsloth import FastLanguageModel
import torch
import json
from datasets import load_from_disk, load_dataset
from trl import SFTTrainer
from transformers import TrainingArguments
from unsloth import is_bfloat16_supported

max_seq_length = 2048
dtype = None
load_in_4bit = True
fourbit_models = [
            "/data/fin-advisor/model/Qwen2-7B-Instruct",
            ]
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "/data/fin-advisor/model/Qwen2-7B-Instruct",
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = load_in_4bit,
)

model = FastLanguageModel.get_peft_model(
    model,
    r = 16,
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
    "gate_proj", "up_proj", "down_proj",],
    lora_alpha = 16,
    lora_dropout = 0,
    bias = "none",
    use_gradient_checkpointing = "unsloth",
    random_state = 3407,
    use_rslora = False,
    loftq_config = None,
)


zalopay_prompt = """
### Instruction:
You are Toro, a helpful Zalopay assistant.
Your answers are correct, high-quality, and written by Zalopay customer support.
You only answer to question about Zalopay to provide helpful answer to user who want to learn more about Zalopay and services provided by Zalopay
### Input:
{}
### Response:
{}
"""
EOS_TOKEN = tokenizer.eos_token # Must add EOS_TOKEN
def formatting_prompts_func(samples):
    questions = samples["question"]
    answers       = samples["answer"]
    texts = []
    for answer, question in zip(questions, answers):
        text = zalopay_prompt.format(answer, question) + EOS_TOKEN
        texts.append(text)
    return { "text" : texts, }



def train(dataset):
    trainer = SFTTrainer(
        model = model,
        tokenizer = tokenizer,
        train_dataset = dataset,
        dataset_text_field = "text",
        max_seq_length = max_seq_length,
        dataset_num_proc = 2,
        args = TrainingArguments(
            per_device_train_batch_size = 2,
            gradient_accumulation_steps = 8,
            # Use num_train_epochs = 1, warmup_ratio for full training runs!
            warmup_steps = 20,
            max_steps = 120,
            learning_rate = 5e-5,
            fp16 = not is_bfloat16_supported(),
            bf16 = is_bfloat16_supported(),
            logging_steps = 1,
            optim = "adamw_8bit",
            weight_decay = 0.01,
            lr_scheduler_type = "linear",
            seed = 3407,
            output_dir = "outputs",
        ),
    )
    trainer_stats = trainer.train()




if __name__ == '__main__':
    samples = []
    with open("data.jsonl", encoding='UTF-8', mode="r") as f:
        for line in f:
            samples.append(json.loads(line))

    dataset = load_dataset("json", data_files="./data.jsonl", split="train")
    dataset = dataset.map(formatting_prompts_func, batched=True)

    # dataset = formatting_prompts_func(samples)
    train(dataset)