from modelscope import snapshot_download
model_dir = snapshot_download("vivo-ai/BlueLM-7B-Chat", revision="v1.0.2")

import torch
from modelscope import AutoModelForCausalLM, AutoTokenizer, snapshot_download
model_dir = snapshot_download("vivo-ai/BlueLM-7B-Chat", revision="v1.0.2")
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True, use_fast=False)
model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="cuda:0", torch_dtype=torch.bfloat16, trust_remote_code=True)
model = model.eval()
inputs = tokenizer("[|Human|]:三国演义的作者是谁？[|AI|]:", return_tensors="pt")
inputs = inputs.to("cuda:0")
pred = model.generate(**inputs, max_new_tokens=64, repetition_penalty=1.1)
print(tokenizer.decode(pred.cpu()[0], skip_special_tokens=True))

# 设置pip全局镜像和安装相关的python包
pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/
git clone https://github.com/modelscope/swift.git
cd swift
pip install .[llm]
# 下面的脚本需要在此目录下执行
cd examples/pytorch/llm
# 如果你想要使用deepspeed.
pip install deepspeed -U
# 如果你想要使用基于bnb的qlora训练.
pip install bitsandbytes -U

# Experimental environment: A10, 3090
# 17GB GPU memory
PYTHONPATH=../../.. \
CUDA_VISIBLE_DEVICES=0 \
python llm_sft.py \
    --model_id_or_path vivo-ai/BlueLM-7B-Chat \
    --model_revision master \
    --sft_type lora \
    --tuner_backend swift \
    --template_type bluelm \
    --dtype bf16 \
    --output_dir output \
    --dataset blossom-math-zh \
    --train_dataset_sample -1 \
    --num_train_epochs 1 \
    --max_length 2048 \
    --check_dataset_strategy warning \
    --lora_rank 8 \
    --lora_alpha 32 \
    --lora_dropout_p 0.05 \
    --lora_target_modules AUTO \
    --gradient_checkpointing true \
    --batch_size 1 \
    --weight_decay 0.01 \
    --learning_rate 1e-4 \
    --gradient_accumulation_steps 16 \
    --max_grad_norm 0.5 \
    --warmup_ratio 0.03 \
    --eval_steps 100 \
    --save_steps 100 \
    --save_total_limit 2 \
    --logging_steps 10 \
    --push_to_hub false \
    --hub_model_id bluelm-7b-chat-lora \
    --hub_private_repo true \
    --hub_token 'your-sdk-token' \

# Experimental environment: A10, 3090
PYTHONPATH=../../.. \
CUDA_VISIBLE_DEVICES=0 \
python llm_infer.py \
    --ckpt_dir "output/bluelm-7b-chat/vx_xxx/checkpoint-xxx" \
    --load_args_from_ckpt_dir true \
    --eval_human false \
    --max_length 2048 \
    --max_new_tokens 2048 \
    --temperature 0.9 \
    --top_k 20 \
    --top_p 0.9 \
    --repetition_penalty 1.05 \
    --do_sample true \
    --merge_lora_and_save false \

[PROMPT]<s> [|Human|]: 学校需新添28套课桌椅，共付款1820元，椅子每把17元，课桌每张多少元？[|AI|]:[OUTPUT]设课桌每张的价格为x元。
根据题意，椅子总价格为28 * 17 = 476元。
课桌的总价格为1820 - 476 = 1344元。
设课桌每张的价格为x元，则28 * x = 1344。
解方程得：x = 1344 / 28 = 48。
所以，课桌每张的价格为48元。
Answer:  48</s>
[LABELS]假设课桌的价格为x元每张。
根据题意可得方程：
28(17+x) = 1820
化简得：
476 + 28x = 1820
移项得：
28x = 1820 - 476
化简得：
28x = 1344
将等式两边同时除以28，得：
x = 1344 / 28
化简得：
x = 48
所以，每张课桌的价格为48元。
Answer: 48
--------------------------------------------------------------------------------
[PROMPT]<s> [|Human|]: 一种小家电，打0.8后的价格是160元，这种小家电原价多少元． [|AI|]:[OUTPUT]根据题意，打0.8后的价格为160元，我们可以表示为：
原价 × 0.8 = 160元
移项得：
原价 = 160元 / 0.8
计算可得：
原价 = 200元
因此，这种小家电的原价为200元。
Answer:  200</s>
[LABELS]设小家电的原价为x元。
根据题意，打0.8折扣后的价格为0.8x元。
根据题意，0.8x = 160。
解这个方程可以得到x = 200。
所以，这种小家电的原价为200元。
Answer: 200