import json
from typing import Any

import torch

from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.quantizers import HfQuantizer, register_quantization_config, register_quantizer
from transformers.utils.quantization_config import QuantizationConfigMixin


@register_quantization_config("custom")
class CustomConfig(QuantizationConfigMixin):
    def __init__(self):
        self.quant_method = "custom"
        self.bits = 8

    def to_dict(self) -> dict[str, Any]:
        output = {
            "num_bits": self.bits,
        }
        return output

    def __repr__(self):
        config_dict = self.to_dict()
        return f"{self.__class__.__name__} {json.dumps(config_dict, indent=2, sort_keys=True)}\n"

    def to_diff_dict(self) -> dict[str, Any]:
        config_dict = self.to_dict()

        default_config_dict = CustomConfig().to_dict()

        serializable_config_dict = {}

        for key, value in config_dict.items():
            if value != default_config_dict[key]:
                serializable_config_dict[key] = value

        return serializable_config_dict


@register_quantizer("custom")
class CustomQuantizer(HfQuantizer):
    def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs):
        super().__init__(quantization_config, **kwargs)
        self.quantization_config = quantization_config
        self.scale_map = {}
        self.device = kwargs.get("device", "cuda" if torch.cuda.is_available() else "cpu")
        self.torch_dtype = kwargs.get("torch_dtype", torch.float32)

    def _process_model_before_weight_loading(self, model, **kwargs):
        return True

    def _process_model_after_weight_loading(self, model, **kwargs):
        return True

    def is_serializable(self) -> bool:
        return True

    def is_trainable(self) -> bool:
        return False


model_8bit = AutoModelForCausalLM.from_pretrained(
    "facebook/opt-350m", quantization_config=CustomConfig(), torch_dtype="auto"
)

tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
input_text = "once there is"
inputs = tokenizer(input_text, return_tensors="pt")
output = model_8bit.generate(
    **inputs,
    max_length=100,
    num_return_sequences=1,
    no_repeat_ngram_size=2,
)
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)

print(generated_text)