import sys

from transformers import AutoTokenizer, AutoModelForCausalLM
from llmcompressor.transformers import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier

MODEL_ID = sys.argv[1]
SAVE_DIR = sys.argv[2]

# Load the model
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto", torch_dtype="auto", local_files_only=True, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, local_files_only=True, trust_remote_code=True)

# Configure the simple PTQ quantization
recipe = QuantizationModifier(targets="Linear", scheme="FP8_DYNAMIC", ignore=["lm_head"])

# Apply the quantization algorithm.
oneshot(model=model, recipe=recipe, trust_remote_code_model=True)

# Save the model
model.save_pretrained(SAVE_DIR)
tokenizer.save_pretrained(SAVE_DIR)