|
|
from datasets import load_dataset |
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoTokenizer |
|
|
import torch |
|
|
|
|
|
|
|
|
|
|
|
dataset_name = 'AlexanderDoria/novel17_test' |
|
|
dataset = load_dataset(dataset_name, split="train") |
|
|
demo.launch() |
|
|
|
|
|
|
|
|
|
|
|
model_name = "TinyPixel/Llama-2-7B-bf16-sharded" |
|
|
|
|
|
bnb_config = BitsAndBytesConfig( |
|
|
load_in_4bit=True, |
|
|
bnb_4bit_quant_type="nf4", |
|
|
bnb_4bit_compute_dtype=torch.float16, |
|
|
) |
|
|
|
|
|
model = AutoModelForCausalLM.from_pretrained( |
|
|
model_name, |
|
|
quantization_config=bnb_config, |
|
|
trust_remote_code=True |
|
|
) |
|
|
model.config.use_cache = False |
