导包

1
2
3
4
5
import torch

from peft import prepare_model_for_kbit_training
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from transformers import pipeline

加载模型

我们使用bitsandbytes加载模型。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
MODEL_NAME = 'model/phi-2/'
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map="auto",
    trust_remote_code=True,
    quantization_config=bnb_config
)
model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=False)

这里使用了本地的模型'model/phi-2/'，也可以加载线上模型。

加载分词器

1
2
3
# tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer.pad_token = tokenizer.eos_token

加载pipeline

1
2
3
4
5
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer
)

prompt生成文本

1
2
3
4
5
6
7
8
9
10
11
12
13
prompt = "As a data scientist, can you explain the concept of regularization in machine learning?"

sequences = pipe(
    prompt,
    do_sample=True,
    max_new_tokens=100,
    temperature=0.7,
    top_k=50,
    top_p=0.95,
    num_return_sequences=1,
)

print(sequences[0]['generated_text'])