```python import time from datasets import load_dataset from transformers import AutoTokenizer from rich import print from llmcompressor.modifiers.quantization import GPTQModifier from llmcompressor.transformers import SparseAutoModelForCausalLM, oneshot # Select model and load it. MODEL_ID = "Unbabel/TowerInstruct-7B-v0.1" model = SparseAutoModelForCausalLM.from_pretrained( MODEL_ID, device_map="auto", torch_dtype="auto", ) tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) # Select calibration dataset. DATASET_ID = "neuralmagic/LLM_compression_calibration" DATASET_SPLIT = "train" # Select number of samples. 512 samples is a good place to start. # Increasing the number of samples can improve accuracy. NUM_CALIBRATION_SAMPLES = 756 MAX_SEQUENCE_LENGTH = 2048 # Load dataset and preprocess. ds = load_dataset(DATASET_ID, split=DATASET_SPLIT) ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES)) def preprocess(example): return { "text": tokenizer.apply_chat_template( example["messages"], tokenize=False, ) } ds = ds.map(preprocess) # Tokenize inputs. def tokenize(sample): return tokenizer( sample["text"], padding=False, max_length=MAX_SEQUENCE_LENGTH, truncation=True, add_special_tokens=False, ) ds = ds.map(tokenize, remove_columns=ds.column_names) # Configure the quantization algorithm to run. # * quantize the weights to 4 bit with GPTQ with a group size 128 # Note: to reduce GPU memory use `sequential_update=False` recipe = GPTQModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"]) print(recipe) # Apply algorithms. oneshot( model=model, dataset=ds, recipe=recipe, max_seq_length=MAX_SEQUENCE_LENGTH, num_calibration_samples=NUM_CALIBRATION_SAMPLES, ) # Confirm generations of the quantized model look sane and measure generation time. print("\n\n") print("========== SAMPLE GENERATION ==============") input_text = "Translate the following text from Portuguese into English.\nPortuguese: Um grupo de investigadores lançou um novo modelo para tarefas relacionadas com tradução.\nEnglish:" input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda") start_time = time.time() output = model.generate(input_ids, max_new_tokens=100) end_time = time.time() generation_time = end_time - start_time print(tokenizer.decode(output[0])) print(f"Generation time: {generation_time:.2f} seconds") print("==========================================\n\n") # Save to disk compressed. SAVE_DIR = MODEL_ID.split("/")[1] + "-quantized.w4a16" model.save_pretrained(SAVE_DIR, save_compressed=True) tokenizer.save_pretrained(SAVE_DIR) ```