from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch

### Config Quantização
quantization = BitsAndBytesConfig (
    load_in_4bit = True,
    bnb_4bit_quant_type = 'nf4',  #Tipo de Quantização
    bnb_4bit_use_double_quant = True, #Double Quantization
    bnb_4bit_compute_dtype = 'float16', #Tipo de precisão usada nos cálculos durante a inferência
)


# Caminho do modelo caso seja self hosted
PATH = r""
# Nome do modelo HF
MODEL = "amalia-llmAMALIA-9B-0626-SFT"

modelo = AutoModelForCausalLM.from_pretrained (PATH, device_map = "auto", quantization_config = quantization)
tokenizer = AutoTokenizer.from_pretrained (PATH)

modelo.save_pretrained (r"")
tokenizer.save_pretrained (r"")