Views
No views yet
1import sys
2import random
3import string
4from transformers import AutoProcessor, AutoModelForCausalLM
5from datasets import load_dataset
6from llmcompressor import oneshot
7from llmcompressor.modifiers.quantization import GPTQModifier
8import torch
9
10MODEL_ID = sys.argv[1]
11NUM_CALIBRATION_SAMPLES=1024
12MAX_SEQUENCE_LENGTH=2048
13
14model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype="auto")
15processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
16
17dampening_frac=0.07
18
19ds = load_dataset("Helsinki-NLP/opus-100", "en-ja", split="train[:1024]")
20def preprocess_function(example):
21 en = example["translation"].get("en", "")
22 ja = example["translation"].get("ja", "")
23 rid = ''.join(random.choices(string.ascii_lowercase + string.digits, k=8))
24 messages = [
25 {"role": "system", "content": "Translate the user's text to Japanese.\nOutput format:XML\nProvide the final translation immediately without any other text."},
26 {"role": "user", "content": f'<seg id="{rid}" type="calib"><field key="content">{en}</field></seg>'},
27 {"role": "assistant", "content": f'<seg id="{rid}" type="calib"><field key="content">{ja}</field></seg>'},
28 ]
29 return processor.apply_chat_template(
30 messages,
31 return_tensors="pt",
32 padding=False,
33 truncation=True,
34 max_length=MAX_SEQUENCE_LENGTH,
35 tokenize=True,
36 add_special_tokens=False,
37 return_dict=True,
38 add_generation_prompt=False,
39 )
40ds = ds.map(preprocess_function, batched=False, remove_columns=ds.column_names)
41
42def data_collator(batch):
43 assert len(batch) == 1
44 return {
45 key: (
46 torch.tensor(value)
47 if key != "pixel_values"
48 else torch.tensor(value, dtype=torch.bfloat16).squeeze(0)
49 )
50 for key, value in batch[0].items()
51 }
52
53recipe = [
54 GPTQModifier(
55 targets="Linear",
56 scheme="W4A16",
57 ignore=["lm_head"],
58 dampening_frac=dampening_frac,
59 )
60]
61
62SAVE_DIR = sys.argv[2]
63oneshot(
64 model=model,
65 processor=processor,
66 recipe=recipe,
67 dataset=ds,
68 max_seq_length=MAX_SEQUENCE_LENGTH,
69 num_calibration_samples=NUM_CALIBRATION_SAMPLES,
70 data_collator=data_collator,
71 sequential_targets=["Gemma3DecoderLayer"],
72 tie_word_embeddings=True,
73)
74
75model.save_pretrained(SAVE_DIR, save_compressed=True)
76processor.save_pretrained(SAVE_DIR)